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 2015/07/23 17:41:09 UTC

[1/6] flink git commit: [FLINK-2332] [runtime] Adds leader session IDs and registration session IDs to JobManager and TaskManager messages.

Repository: flink
Updated Branches:
  refs/heads/master fa78be6df -> 222245428


http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationMasterActor.scala
----------------------------------------------------------------------
diff --git a/flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationMasterActor.scala b/flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationMasterActor.scala
index 3fb5e30..e99f8d2 100644
--- a/flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationMasterActor.scala
+++ b/flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationMasterActor.scala
@@ -25,7 +25,7 @@ import java.util.Collections
 import akka.actor.ActorRef
 import org.apache.flink.api.common.JobID
 import org.apache.flink.configuration.ConfigConstants
-import org.apache.flink.runtime.ActorLogMessages
+import org.apache.flink.runtime.FlinkActor
 import org.apache.flink.runtime.jobgraph.JobStatus
 import org.apache.flink.runtime.jobmanager.JobManager
 import org.apache.flink.runtime.messages.JobManagerMessages.{CurrentJobStatus, JobNotFound, RequestJobStatus}
@@ -52,7 +52,7 @@ import scala.language.postfixOps
 import scala.util.Try
 
 
-trait ApplicationMasterActor extends ActorLogMessages {
+trait ApplicationMasterActor extends FlinkActor {
   that: JobManager =>
 
   import context._
@@ -92,18 +92,17 @@ trait ApplicationMasterActor extends ActorLogMessages {
   var allocatedContainersList: mutable.MutableList[Container] = new mutable.MutableList[Container]
   var runningContainersList: mutable.MutableList[Container] = new mutable.MutableList[Container]
 
-
-  abstract override def receiveWithLogMessages: Receive = {
-    receiveYarnMessages orElse super.receiveWithLogMessages
+  abstract override def handleMessage: Receive = {
+    handleYarnMessage orElse super.handleMessage
   }
 
-  def receiveYarnMessages: Receive = {
+  def handleYarnMessage: Receive = {
     case StopYarnSession(status, diag) =>
       log.info(s"Stopping YARN JobManager with status $status and diagnostic $diag.")
 
       instanceManager.getAllRegisteredInstances.asScala foreach {
         instance =>
-          instance.getInstanceGateway.tell(StopYarnSession(status, diag))
+          instance.getActorGateway.tell(StopYarnSession(status, diag))
       }
 
       rmClientOption foreach {
@@ -128,7 +127,7 @@ trait ApplicationMasterActor extends ActorLogMessages {
 
       nmClientOption = None
       messageListener foreach {
-          _ ! JobManagerStopped
+          _ ! decorateMessage(JobManagerStopped)
       }
 
       context.system.shutdown()
@@ -136,7 +135,7 @@ trait ApplicationMasterActor extends ActorLogMessages {
     case RegisterClient(client) =>
       log.info(s"Register ${client.path} as client.")
       messageListener = Some(client)
-      sender ! Acknowledge
+      sender ! decorateMessage(Acknowledge)
 
     case UnregisterClient =>
       messageListener = None
@@ -145,12 +144,15 @@ trait ApplicationMasterActor extends ActorLogMessages {
       val jobId = msg.jobId
       log.info(s"ApplicatonMaster will shut down YARN session when job $jobId has finished.")
       stopWhenJobFinished = jobId
-      sender() ! Acknowledge
+      sender() ! decorateMessage(Acknowledge)
 
 
     case PollYarnClusterStatus =>
-      sender() ! new FlinkYarnClusterStatus(instanceManager.getNumberOfRegisteredTaskManagers,
-        instanceManager.getTotalNumberOfSlots)
+      sender() ! decorateMessage(
+        new FlinkYarnClusterStatus(
+          instanceManager.getNumberOfRegisteredTaskManagers,
+          instanceManager.getTotalNumberOfSlots)
+      )
 
     case StartYarnSession(conf, actorSystemPort, webServerPort) =>
       startYarnSession(conf, actorSystemPort, webServerPort)
@@ -173,11 +175,17 @@ trait ApplicationMasterActor extends ActorLogMessages {
             log.info(s"Job with ID ${jobStatus.jobID} is in terminal state ${jobStatus.status}. " +
               s"Shutting down YARN session")
             if (jobStatus.status == JobStatus.FINISHED) {
-              self ! StopYarnSession(FinalApplicationStatus.SUCCEEDED,
-                s"The monitored job with ID ${jobStatus.jobID} has finished.")
+              self ! decorateMessage(
+                StopYarnSession(
+                  FinalApplicationStatus.SUCCEEDED,
+                  s"The monitored job with ID ${jobStatus.jobID} has finished.")
+              )
             } else {
-              self ! StopYarnSession(FinalApplicationStatus.FAILED,
-                s"The monitored job with ID ${jobStatus.jobID} has failed to complete.")
+              self ! decorateMessage(
+                StopYarnSession(
+                  FinalApplicationStatus.FAILED,
+                  s"The monitored job with ID ${jobStatus.jobID} has failed to complete.")
+              )
             }
           } else {
             log.debug(s"Monitored job with ID ${jobStatus.jobID} is in state ${jobStatus.status}")
@@ -188,7 +196,7 @@ trait ApplicationMasterActor extends ActorLogMessages {
     case HeartbeatWithYarn =>
       // piggyback on the YARN heartbeat to check if the job has finished
       if(stopWhenJobFinished != null) {
-        self ! RequestJobStatus(stopWhenJobFinished)
+        self ! decorateMessage(RequestJobStatus(stopWhenJobFinished))
       }
       rmClientOption match {
         case Some(rmClient) =>
@@ -226,8 +234,10 @@ trait ApplicationMasterActor extends ActorLogMessages {
                   case _ => ""
                 }
                 messageListener foreach {
-                  _ ! YarnMessage(s"Diagnostics for containerID=${status.getContainerId} in " +
+                  _ ! decorateMessage(
+                    YarnMessage(s"Diagnostics for containerID=${status.getContainerId} in " +
                     s"state=${status.getState}.\n${status.getDiagnostics} $detail")
+                  )
                 }
               }
               // return
@@ -281,7 +291,7 @@ trait ApplicationMasterActor extends ActorLogMessages {
                             containersLaunched += 1
                             runningContainersList += container
                             messageListener foreach {
-                              _ ! YarnMessage(message)
+                              _ ! decorateMessage(YarnMessage(message))
                             }
                           } catch {
                             case e: YarnException =>
@@ -290,13 +300,19 @@ trait ApplicationMasterActor extends ActorLogMessages {
                         }
                         case None =>
                           log.error("The ContainerLaunchContext was not set.")
-                          self ! StopYarnSession(FinalApplicationStatus.FAILED,
-                            "Fatal error in AM: The ContainerLaunchContext was not set.")
+                          self ! decorateMessage(
+                            StopYarnSession(
+                              FinalApplicationStatus.FAILED,
+                              "Fatal error in AM: The ContainerLaunchContext was not set.")
+                          )
                       }
                     case None =>
                       log.error("The NMClient was not set.")
-                      self ! StopYarnSession(FinalApplicationStatus.FAILED,
-                        "Fatal error in AM: The NMClient was not set.")
+                      self ! decorateMessage(
+                        StopYarnSession(
+                          FinalApplicationStatus.FAILED,
+                          "Fatal error in AM: The NMClient was not set.")
+                      )
                   }
                   // dropping condition
                   true
@@ -343,21 +359,30 @@ trait ApplicationMasterActor extends ActorLogMessages {
               s"the '${ConfigConstants.YARN_MAX_FAILED_CONTAINERS}' configuration " +
               s"setting. By default its the number of requested containers"
             log.error(msg)
-            self ! StopYarnSession(FinalApplicationStatus.FAILED, msg)
+            self ! decorateMessage(StopYarnSession(FinalApplicationStatus.FAILED, msg))
           }
 
           // schedule next heartbeat:
           if (runningContainers < numTaskManager) {
             // we don't have the requested number of containers. Do fast polling
-            context.system.scheduler.scheduleOnce(FAST_YARN_HEARTBEAT_DELAY, self,HeartbeatWithYarn)
+            context.system.scheduler.scheduleOnce(
+              FAST_YARN_HEARTBEAT_DELAY,
+              self,
+              decorateMessage(HeartbeatWithYarn))
           } else {
             // everything is good, slow down polling
-            context.system.scheduler.scheduleOnce(YARN_HEARTBEAT_DELAY, self, HeartbeatWithYarn)
+            context.system.scheduler.scheduleOnce(
+              YARN_HEARTBEAT_DELAY,
+              self,
+              decorateMessage(HeartbeatWithYarn))
           }
         case None =>
           log.error("The AMRMClient was not set.")
-          self ! StopYarnSession(FinalApplicationStatus.FAILED, "Fatal error in AM: AMRMClient " +
-            "was not set")
+          self ! decorateMessage(
+            StopYarnSession(
+              FinalApplicationStatus.FAILED,
+              "Fatal error in AM: AMRMClient was not set")
+          )
       }
       log.debug(s"Processed Heartbeat with RMClient. Running containers $runningContainers, " +
         s"failed containers $failedContainers, " +
@@ -371,9 +396,12 @@ trait ApplicationMasterActor extends ActorLogMessages {
     allocatedContainersList map { runningCont => runningCont.getId}
   }
 
-  private def startYarnSession(conf: Configuration,
-                               actorSystemPort: Int,
-                               webServerPort: Int): Unit = {
+  private def startYarnSession(
+      conf: Configuration,
+      actorSystemPort: Int,
+      webServerPort: Int)
+    : Unit = {
+
     Try {
       log.info("Start yarn session.")
       memoryPerTaskManager = env.get(FlinkYarnClient.ENV_TM_MEMORY).toInt
@@ -383,8 +411,10 @@ trait ApplicationMasterActor extends ActorLogMessages {
       require(applicationMasterHost != null, s"Application master (${Environment.NM_HOST} not set.")
 
       val yarnExpiryInterval: FiniteDuration = FiniteDuration(
-        conf.getInt(YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS,
-          YarnConfiguration.DEFAULT_RM_AM_EXPIRY_INTERVAL_MS), MILLISECONDS)
+        conf.getInt(
+          YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS,
+          YarnConfiguration.DEFAULT_RM_AM_EXPIRY_INTERVAL_MS),
+        MILLISECONDS)
 
       if(YARN_HEARTBEAT_DELAY.gteq(yarnExpiryInterval)) {
         log.warn(s"The heartbeat interval of the Flink Application master " +
@@ -468,16 +498,30 @@ trait ApplicationMasterActor extends ActorLogMessages {
       failedContainers = 0
 
       val hs = ApplicationMaster.hasStreamingMode(env)
-      containerLaunchContext = Some(createContainerLaunchContext(heapLimit, hasLogback, hasLog4j,
-        yarnClientUsername, conf, taskManagerLocalResources, hs))
-
-
-      context.system.scheduler.scheduleOnce(FAST_YARN_HEARTBEAT_DELAY, self, HeartbeatWithYarn)
+      containerLaunchContext = Some(
+        createContainerLaunchContext(
+          heapLimit,
+          hasLogback,
+          hasLog4j,
+          yarnClientUsername,
+          conf,
+          taskManagerLocalResources,
+          hs)
+      )
+
+
+      context.system.scheduler.scheduleOnce(
+        FAST_YARN_HEARTBEAT_DELAY,
+        self,
+        decorateMessage(HeartbeatWithYarn))
     } recover {
       case t: Throwable =>
         log.error("Could not start yarn session.", t)
-        self ! StopYarnSession(FinalApplicationStatus.FAILED,
-          s"ApplicationMaster failed while starting. Exception Message: ${t.getMessage}")
+        self ! decorateMessage(
+          StopYarnSession(
+            FinalApplicationStatus.FAILED,
+            s"ApplicationMaster failed while starting. Exception Message: ${t.getMessage}")
+        )
     }
   }
 
@@ -505,11 +549,15 @@ trait ApplicationMasterActor extends ActorLogMessages {
     new ContainerRequest(capability, null, null, priority)
   }
 
-  private def createContainerLaunchContext(heapLimit: Int, hasLogback: Boolean, hasLog4j: Boolean,
-                                   yarnClientUsername: String, yarnConf: Configuration,
-                                   taskManagerLocalResources: Map[String, LocalResource],
-                                   streamingMode: Boolean):
-  ContainerLaunchContext = {
+  private def createContainerLaunchContext(
+      heapLimit: Int,
+      hasLogback: Boolean,
+      hasLog4j: Boolean,
+      yarnClientUsername: String,
+      yarnConf: Configuration,
+      taskManagerLocalResources: Map[String, LocalResource],
+      streamingMode: Boolean)
+    : ContainerLaunchContext = {
     log.info("Create container launch context.")
     val ctx = Records.newRecord(classOf[ContainerLaunchContext])
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-yarn/src/main/scala/org/apache/flink/yarn/YarnTaskManager.scala
----------------------------------------------------------------------
diff --git a/flink-yarn/src/main/scala/org/apache/flink/yarn/YarnTaskManager.scala b/flink-yarn/src/main/scala/org/apache/flink/yarn/YarnTaskManager.scala
index de82716..5216030 100644
--- a/flink-yarn/src/main/scala/org/apache/flink/yarn/YarnTaskManager.scala
+++ b/flink-yarn/src/main/scala/org/apache/flink/yarn/YarnTaskManager.scala
@@ -25,26 +25,31 @@ import org.apache.flink.runtime.memorymanager.DefaultMemoryManager
 import org.apache.flink.runtime.taskmanager.{NetworkEnvironmentConfiguration, TaskManagerConfiguration, TaskManager}
 import org.apache.flink.yarn.Messages.StopYarnSession
 
-/**
- * An extension of the TaskManager that listens for additional YARN related
- * messages.
- */
-class YarnTaskManager(config: TaskManagerConfiguration,
-                      connectionInfo: InstanceConnectionInfo,
-                      jobManagerAkkaURL: String,
-                      memoryManager: DefaultMemoryManager,
-                      ioManager: IOManager,
-                      network: NetworkEnvironment,
-                      numberOfSlots: Int)
-  extends TaskManager(config, connectionInfo, jobManagerAkkaURL,
-                      memoryManager, ioManager, network, numberOfSlots) {
-
+/** An extension of the TaskManager that listens for additional YARN related
+  * messages.
+  */
+class YarnTaskManager(
+    config: TaskManagerConfiguration,
+    connectionInfo: InstanceConnectionInfo,
+    jobManagerAkkaURL: String,
+    memoryManager: DefaultMemoryManager,
+    ioManager: IOManager,
+    network: NetworkEnvironment,
+    numberOfSlots: Int)
+  extends TaskManager(
+    config,
+    connectionInfo,
+    jobManagerAkkaURL,
+    memoryManager,
+    ioManager,
+    network,
+    numberOfSlots) {
 
-  override def receiveWithLogMessages: Receive = {
-    receiveYarnMessages orElse super.receiveWithLogMessages
+  override def handleMessage: Receive = {
+    handleYarnMessages orElse super.handleMessage
   }
 
-  def receiveYarnMessages: Receive = {
+  def handleYarnMessages: Receive = {
     case StopYarnSession(status, diagnostics) =>
       log.info(s"Stopping YARN TaskManager with final application status $status " +
         s"and diagnostics: $diagnostics")


[6/6] flink git commit: [FLINK-2332] [runtime] Adds leader session IDs and registration session IDs to JobManager and TaskManager messages.

Posted by tr...@apache.org.
[FLINK-2332] [runtime] Adds leader session IDs and registration session IDs to JobManager and TaskManager messages.

Refactors Flink's actor traits to support stackable traits

Refactored Flink's actors to use a factory method to generate messages

Replaced ActorRef with InstanceGateway in Task, RuntimeEnvironment and ExecutionGraph

Add comments

Add test cases for registration session ID and leader session ID

Adds IT case to check that a CancelMessage with the wrong leader session ID is discarded

Corrected order of visibility and abstract modifiers. Removed lazy log member from FlinkActor. Made RequiresLeaderSessionID a Java interface.

This closes #917.


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

Branch: refs/heads/master
Commit: 222245428518d0e1b843947a6184b4a803a78ad5
Parents: fa78be6
Author: Till Rohrmann <tr...@apache.org>
Authored: Wed Jul 8 15:37:50 2015 +0200
Committer: Till Rohrmann <tr...@apache.org>
Committed: Thu Jul 23 17:39:57 2015 +0200

----------------------------------------------------------------------
 .../org/apache/flink/client/CliFrontend.java    |  28 +-
 .../org/apache/flink/client/program/Client.java |  32 +-
 .../flink/client/CliFrontendListCancelTest.java |  84 ++-
 .../apache/flink/client/program/ClientTest.java |  43 +-
 .../flink/runtime/akka/FlinkUntypedActor.java   | 140 +++++
 .../checkpoint/CheckpointCoordinator.java       |  40 +-
 .../CheckpointCoordinatorDeActivator.java       |  25 +-
 .../apache/flink/runtime/client/JobClient.java  |  69 ++-
 .../flink/runtime/client/JobClientActor.java    |  62 +-
 .../flink/runtime/executiongraph/Execution.java |  10 +-
 .../runtime/executiongraph/ExecutionGraph.java  |  55 +-
 .../runtime/executiongraph/ExecutionVertex.java |   4 +-
 .../flink/runtime/instance/ActorGateway.java    | 107 ++++
 .../runtime/instance/AkkaActorGateway.java      | 157 +++++
 .../runtime/instance/AkkaInstanceGateway.java   | 111 ----
 .../apache/flink/runtime/instance/Instance.java |  14 +-
 .../flink/runtime/instance/InstanceGateway.java |  82 ---
 .../flink/runtime/instance/InstanceManager.java |  24 +-
 .../runtime/io/network/NetworkEnvironment.java  |  47 +-
 .../jobmanager/web/SetupInfoServlet.java        |   2 +-
 .../messages/RequiresLeaderSessionID.java       |  25 +
 .../runtime/taskmanager/RuntimeEnvironment.java |  40 +-
 .../apache/flink/runtime/taskmanager/Task.java  |  39 +-
 .../taskmanager/TaskInputSplitProvider.java     |  24 +-
 .../apache/flink/runtime/ActorLogMessages.scala |  52 --
 .../flink/runtime/ActorSynchronousLogging.scala |  31 -
 .../org/apache/flink/runtime/FlinkActor.scala   |  48 ++
 .../runtime/LeaderSessionMessageDecorator.scala |  46 ++
 .../flink/runtime/LeaderSessionMessages.scala   |  73 +++
 .../org/apache/flink/runtime/LogMessages.scala  |  48 ++
 .../apache/flink/runtime/MessageDecorator.scala |  32 ++
 .../apache/flink/runtime/akka/AkkaUtils.scala   |   6 +-
 .../flink/runtime/jobmanager/JobManager.scala   | 337 +++++++----
 .../runtime/jobmanager/MemoryArchivist.scala    |  33 +-
 .../messages/ExecutionGraphMessages.scala       |  25 +-
 .../runtime/messages/JobManagerMessages.scala   |  44 +-
 .../flink/runtime/messages/Messages.scala       |   2 +-
 .../runtime/messages/RegistrationMessages.scala |  42 +-
 .../flink/runtime/messages/TaskMessages.scala   |  48 +-
 .../runtime/minicluster/FlinkMiniCluster.scala  |  19 +-
 .../flink/runtime/taskmanager/TaskManager.scala | 566 ++++++++++++-------
 .../runtime/akka/FlinkUntypedActorTest.java     | 146 +++++
 .../checkpoint/CoordinatorShutdownTest.java     |  22 +-
 .../ExecutionGraphDeploymentTest.java           |   4 +-
 .../executiongraph/ExecutionGraphTestUtils.java |  17 +-
 .../ExecutionStateProgressTest.java             |   2 +-
 .../ExecutionVertexCancelTest.java              |  40 +-
 .../ExecutionVertexDeploymentTest.java          |  14 +-
 .../ExecutionVertexSchedulingTest.java          |   8 +-
 .../executiongraph/LocalInputSplitsTest.java    |   2 +-
 .../TerminalStateDeadlockTest.java              |   4 +-
 .../VertexLocationConstraintTest.java           |   6 +-
 .../instance/BaseTestingActorGateway.java       | 120 ++++
 .../instance/BaseTestingInstanceGateway.java    |  94 ---
 .../runtime/instance/DummyActorGateway.java     |  68 +++
 .../runtime/instance/DummyInstanceGateway.java  |  57 --
 .../runtime/instance/InstanceManagerTest.java   |  23 +-
 .../flink/runtime/instance/InstanceTest.java    |   6 +-
 .../flink/runtime/instance/SimpleSlotTest.java  |   2 +-
 .../io/network/NetworkEnvironmentTest.java      |  16 +-
 .../PartialConsumePipelinedResultTest.java      |   2 +-
 .../runtime/jobmanager/JobManagerTest.java      |  46 +-
 .../flink/runtime/jobmanager/JobSubmitTest.java |  33 +-
 .../scheduler/SchedulerTestUtils.java           |   4 +-
 .../taskmanager/ForwardingActorGateway.java     |  47 ++
 .../runtime/taskmanager/TaskAsyncCallTest.java  |  35 +-
 .../runtime/taskmanager/TaskCancelTest.java     |  11 +-
 .../taskmanager/TaskInputSplitProviderTest.java |  44 +-
 .../TaskManagerRegistrationTest.java            |  95 +++-
 .../runtime/taskmanager/TaskManagerTest.java    | 170 ++++--
 .../flink/runtime/taskmanager/TaskTest.java     |  93 ++-
 .../flink/runtime/akka/FlinkActorTest.scala     | 107 ++++
 .../ExecutionGraphRestartTest.scala             |   6 +-
 .../TaskManagerLossFailsTasksTest.scala         |   6 +-
 .../jobmanager/CoLocationConstraintITCase.scala |  16 +-
 .../runtime/jobmanager/JobManagerITCase.scala   | 108 ++--
 .../jobmanager/JobManagerRegistrationTest.scala |  43 +-
 .../runtime/jobmanager/RecoveryITCase.scala     |  31 +-
 .../runtime/jobmanager/SlotSharingITCase.scala  |  20 +-
 .../TaskManagerFailsWithSlotSharingITCase.scala |  23 +-
 .../testingUtils/ScalaTestingUtils.scala        |  37 ++
 .../testingUtils/TestingJobManager.scala        |  93 +--
 .../TestingJobManagerMessages.scala             |   4 +-
 .../testingUtils/TestingMemoryArchivist.scala   |  14 +-
 .../testingUtils/TestingTaskManager.scala       | 105 ++--
 .../src/test/resources/log4j-test.properties    |   2 +-
 .../src/test/resources/log4j-test.properties    |   2 +-
 .../accumulators/AccumulatorLiveITCase.java     |  13 +-
 .../test/cancelling/CancellingTestBase.java     |   3 +-
 .../LocalFlinkMiniClusterITCase.java            |  19 +-
 .../jobmanager/JobManagerFailsITCase.scala      |  33 +-
 .../JobManagerLeaderSessionIDITSuite.scala      | 112 ++++
 .../taskmanager/TaskManagerFailsITCase.scala    |  41 +-
 .../apache/flink/yarn/ApplicationClient.scala   |  46 +-
 .../flink/yarn/ApplicationMasterActor.scala     | 138 +++--
 .../org/apache/flink/yarn/YarnTaskManager.scala |  39 +-
 96 files changed, 3367 insertions(+), 1641 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-clients/src/main/java/org/apache/flink/client/CliFrontend.java
----------------------------------------------------------------------
diff --git a/flink-clients/src/main/java/org/apache/flink/client/CliFrontend.java b/flink-clients/src/main/java/org/apache/flink/client/CliFrontend.java
index e28f2fd..71b78bd 100644
--- a/flink-clients/src/main/java/org/apache/flink/client/CliFrontend.java
+++ b/flink-clients/src/main/java/org/apache/flink/client/CliFrontend.java
@@ -38,8 +38,6 @@ import java.util.Properties;
 
 import akka.actor.ActorRef;
 import akka.actor.ActorSystem;
-import akka.pattern.Patterns;
-import akka.util.Timeout;
 
 import org.apache.commons.cli.CommandLine;
 import org.apache.flink.api.common.JobSubmissionResult;
@@ -65,6 +63,7 @@ import org.apache.flink.optimizer.plan.FlinkPlan;
 import org.apache.flink.optimizer.plan.OptimizedPlan;
 import org.apache.flink.optimizer.plandump.PlanJSONDumpGenerator;
 import org.apache.flink.runtime.client.JobStatusMessage;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.messages.JobManagerMessages;
 import org.apache.flink.runtime.security.SecurityUtils;
 import org.apache.flink.runtime.util.EnvironmentInformation;
@@ -447,11 +446,12 @@ public class CliFrontend {
 		}
 
 		try {
-			ActorRef jobManager = getJobManager(options);
+			ActorGateway jobManagerGateway = getJobManagerGateway(options);
 
 			LOG.info("Connecting to JobManager to retrieve list of jobs");
-			Future<Object> response = Patterns.ask(jobManager,
-					JobManagerMessages.getRequestRunningJobsStatus(), new Timeout(askTimeout));
+			Future<Object> response = jobManagerGateway.ask(
+					JobManagerMessages.getRequestRunningJobsStatus(),
+					askTimeout);
 
 			Object result;
 			try {
@@ -580,8 +580,8 @@ public class CliFrontend {
 		}
 
 		try {
-			ActorRef jobManager = getJobManager(options);
-			Future<Object> response = Patterns.ask(jobManager, new CancelJob(jobId), new Timeout(askTimeout));
+			ActorGateway jobManager = getJobManagerGateway(options);
+			Future<Object> response = jobManager.ask(new CancelJob(jobId), askTimeout);
 
 			try {
 				Await.result(response, askTimeout);
@@ -722,7 +722,15 @@ public class CliFrontend {
 		return jobManagerAddress;
 	}
 
-	protected ActorRef getJobManager(CommandLineOptions options) throws Exception {
+	/**
+	 * Retrieves the {@link ActorGateway} for the JobManager. The JobManager address is retrieved
+	 * from the provided {@link CommandLineOptions}.
+	 *
+	 * @param options CommandLineOptions specifying the JobManager URL
+	 * @return Gateway to the JobManager
+	 * @throws Exception
+	 */
+	protected ActorGateway getJobManagerGateway(CommandLineOptions options) throws Exception {
 		//TODO: Get ActorRef from YarnCluster if we are in YARN mode.
 
 		InetSocketAddress address = getJobManagerAddress(options);
@@ -745,7 +753,9 @@ public class CliFrontend {
 		LOG.info("Trying to lookup JobManager");
 		ActorRef jmActor = JobManager.getJobManagerRemoteReference(address, actorSystem, lookupTimeout);
 		LOG.info("JobManager is at " + jmActor.path());
-		return jmActor;
+
+		// Retrieve the ActorGateway from the JobManager's ActorRef
+		return JobManager.getJobManagerGateway(jmActor, lookupTimeout);
 	}
 
 	/**

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-clients/src/main/java/org/apache/flink/client/program/Client.java
----------------------------------------------------------------------
diff --git a/flink-clients/src/main/java/org/apache/flink/client/program/Client.java b/flink-clients/src/main/java/org/apache/flink/client/program/Client.java
index 2908f92..50590df 100644
--- a/flink-clients/src/main/java/org/apache/flink/client/program/Client.java
+++ b/flink-clients/src/main/java/org/apache/flink/client/program/Client.java
@@ -52,6 +52,7 @@ import org.apache.flink.core.fs.Path;
 import org.apache.flink.runtime.client.JobClient;
 import org.apache.flink.runtime.client.JobExecutionException;
 import org.apache.flink.runtime.client.SerializedJobExecutionResult;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.runtime.jobmanager.JobManager;
 import org.apache.flink.runtime.messages.JobManagerMessages;
@@ -375,22 +376,35 @@ public class Client {
 			throw new ProgramInvocationException("Could start client actor system.", e);
 		}
 
+		FiniteDuration timeout = AkkaUtils.getTimeout(configuration);
+
 		LOG.info("Looking up JobManager");
-		ActorRef jobManager;
+		ActorGateway jobManagerGateway;
+		ActorRef jobManagerActorRef;
 		try {
-			jobManager = JobManager.getJobManagerRemoteReference(jobManagerAddress, actorSystem, configuration);
-		}
-		catch (IOException e) {
+			jobManagerActorRef = JobManager.getJobManagerRemoteReference(
+					jobManagerAddress,
+					actorSystem,
+					configuration);
+
+
+		} catch (IOException e) {
 			throw new ProgramInvocationException("Failed to resolve JobManager", e);
 		}
-		LOG.info("JobManager runs at " + jobManager.path());
 
-		FiniteDuration timeout = AkkaUtils.getTimeout(configuration);
+		try{
+			jobManagerGateway = JobManager.getJobManagerGateway(jobManagerActorRef, timeout);
+		} catch (Exception e) {
+			throw new ProgramInvocationException("Failed to retrieve the JobManager gateway.", e);
+		}
+
+		LOG.info("JobManager runs at " + jobManagerGateway.path());
+
 		LOG.info("Communication between client and JobManager will have a timeout of " + timeout);
 
 		LOG.info("Checking and uploading JAR files");
 		try {
-			JobClient.uploadJarFiles(jobGraph, jobManager, timeout);
+			JobClient.uploadJarFiles(jobGraph, jobManagerGateway, timeout);
 		}
 		catch (IOException e) {
 			throw new ProgramInvocationException("Could not upload the program's JAR files to the JobManager.", e);
@@ -399,7 +413,7 @@ public class Client {
 		try{
 			if (wait) {
 				SerializedJobExecutionResult result = JobClient.submitJobAndWait(actorSystem, 
-						jobManager, jobGraph, timeout, printStatusDuringExecution);
+						jobManagerGateway, jobGraph, timeout, printStatusDuringExecution);
 				try {
 					return result.toJobExecutionResult(this.userCodeClassLoader);
 				}
@@ -409,7 +423,7 @@ public class Client {
 				}
 			}
 			else {
-				JobClient.submitJobDetached(jobManager, jobGraph, timeout);
+				JobClient.submitJobDetached(jobManagerGateway, jobGraph, timeout);
 				// return a dummy execution result with the JobId
 				return new JobSubmissionResult(jobGraph.getJobID());
 			}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-clients/src/test/java/org/apache/flink/client/CliFrontendListCancelTest.java
----------------------------------------------------------------------
diff --git a/flink-clients/src/test/java/org/apache/flink/client/CliFrontendListCancelTest.java b/flink-clients/src/test/java/org/apache/flink/client/CliFrontendListCancelTest.java
index 3224e0f..fc64503 100644
--- a/flink-clients/src/test/java/org/apache/flink/client/CliFrontendListCancelTest.java
+++ b/flink-clients/src/test/java/org/apache/flink/client/CliFrontendListCancelTest.java
@@ -23,10 +23,16 @@ import akka.testkit.JavaTestKit;
 
 import org.apache.flink.client.cli.CommandLineOptions;
 import org.apache.flink.api.common.JobID;
+import org.apache.flink.runtime.akka.FlinkUntypedActor;
+import org.apache.flink.runtime.instance.AkkaActorGateway;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.messages.JobManagerMessages;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import scala.Option;
+
+import java.util.UUID;
 
 import static org.apache.flink.client.CliFrontendTestUtils.pipeSystemOutToNull;
 import static org.junit.Assert.*;
@@ -77,10 +83,19 @@ public class CliFrontendListCancelTest {
 				JobID jid = new JobID();
 				String jidString = jid.toString();
 
-				final ActorRef jm = actorSystem.actorOf(Props.create(CliJobManager.class, jid));
+				final Option<UUID> leaderSessionID = Option.<UUID>apply(UUID.randomUUID());
+
+				final ActorRef jm = actorSystem.actorOf(Props.create(
+								CliJobManager.class,
+								jid,
+								leaderSessionID
+						)
+				);
+
+				final ActorGateway gateway = new AkkaActorGateway(jm, leaderSessionID);
 				
 				String[] parameters = { jidString };
-				InfoListTestCliFrontend testFrontend = new InfoListTestCliFrontend(jm);
+				InfoListTestCliFrontend testFrontend = new InfoListTestCliFrontend(gateway);
 
 				int retCode = testFrontend.cancel(parameters);
 				assertTrue(retCode == 0);
@@ -91,10 +106,20 @@ public class CliFrontendListCancelTest {
 				JobID jid1 = new JobID();
 				JobID jid2 = new JobID();
 
-				final ActorRef jm = actorSystem.actorOf(Props.create(CliJobManager.class, jid1));
+				final Option<UUID> leaderSessionID = Option.<UUID>apply(UUID.randomUUID());
+
+				final ActorRef jm = actorSystem.actorOf(
+						Props.create(
+								CliJobManager.class,
+								jid1,
+								leaderSessionID
+						)
+				);
+
+				final ActorGateway gateway = new AkkaActorGateway(jm, leaderSessionID);
 
 				String[] parameters = { jid2.toString() };
-				InfoListTestCliFrontend testFrontend = new InfoListTestCliFrontend(jm);
+				InfoListTestCliFrontend testFrontend = new InfoListTestCliFrontend(gateway);
 
 				assertTrue(testFrontend.cancel(parameters) != 0);
 			}
@@ -118,9 +143,17 @@ public class CliFrontendListCancelTest {
 			
 			// test list properly
 			{
-				final ActorRef jm = actorSystem.actorOf(Props.create(CliJobManager.class, (Object)null));
+				final Option<UUID> leaderSessionID = Option.<UUID>apply(UUID.randomUUID());
+				final ActorRef jm = actorSystem.actorOf(
+						Props.create(
+								CliJobManager.class,
+								(Object)null,
+								leaderSessionID
+						)
+				);
+				final ActorGateway gateway = new AkkaActorGateway(jm, leaderSessionID);
 				String[] parameters = {"-r", "-s"};
-				InfoListTestCliFrontend testFrontend = new InfoListTestCliFrontend(jm);
+				InfoListTestCliFrontend testFrontend = new InfoListTestCliFrontend(gateway);
 				int retCode = testFrontend.list(parameters);
 				assertTrue(retCode == 0);
 			}
@@ -135,46 +168,57 @@ public class CliFrontendListCancelTest {
 
 	protected static final class InfoListTestCliFrontend extends CliFrontend {
 
-		private ActorRef jobmanager;
+		private ActorGateway jobManagerGateway;
 
-
-
-		public InfoListTestCliFrontend(ActorRef jobmanager) throws Exception {
+		public InfoListTestCliFrontend(ActorGateway jobManagerGateway) throws Exception {
 			super(CliFrontendTestUtils.getConfigDir());
-			this.jobmanager = jobmanager;
+			this.jobManagerGateway = jobManagerGateway;
 		}
 
 		@Override
-		public ActorRef getJobManager(CommandLineOptions options) {
-			return jobmanager;
+		public ActorGateway getJobManagerGateway(CommandLineOptions options) {
+			return jobManagerGateway;
 		}
 	}
 
-	protected static final class CliJobManager extends UntypedActor{
+	protected static final class CliJobManager extends FlinkUntypedActor {
 		private final JobID jobID;
+		private final Option<UUID> leaderSessionID;
 
-		public CliJobManager(final JobID jobID){
+		public CliJobManager(final JobID jobID, final Option<UUID> leaderSessionID){
 			this.jobID = jobID;
+			this.leaderSessionID = leaderSessionID;
 		}
 
 		@Override
-		public void onReceive(Object message) throws Exception {
+		public void handleMessage(Object message) {
 			if (message instanceof JobManagerMessages.RequestTotalNumberOfSlots$) {
-				getSender().tell(1, getSelf());
+				getSender().tell(decorateMessage(1), getSelf());
 			}
 			else if (message instanceof JobManagerMessages.CancelJob) {
 				JobManagerMessages.CancelJob cancelJob = (JobManagerMessages.CancelJob) message;
 
 				if (jobID != null && jobID.equals(cancelJob.jobID())) {
-					getSender().tell(new Status.Success(new Object()), getSelf());
+					getSender().tell(
+							decorateMessage(new Status.Success(new Object())),
+							getSelf());
 				}
 				else {
-					getSender().tell(new Status.Failure(new Exception("Wrong or no JobID")), getSelf());
+					getSender().tell(
+							decorateMessage(new Status.Failure(new Exception("Wrong or no JobID"))),
+							getSelf());
 				}
 			}
 			else if (message instanceof JobManagerMessages.RequestRunningJobsStatus$) {
-				getSender().tell(new JobManagerMessages.RunningJobsStatus(), getSelf());
+				getSender().tell(
+						decorateMessage(new JobManagerMessages.RunningJobsStatus()),
+						getSelf());
 			}
 		}
+
+		@Override
+		protected Option<UUID> getLeaderSessionID() {
+			return leaderSessionID;
+		}
 	}
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-clients/src/test/java/org/apache/flink/client/program/ClientTest.java
----------------------------------------------------------------------
diff --git a/flink-clients/src/test/java/org/apache/flink/client/program/ClientTest.java b/flink-clients/src/test/java/org/apache/flink/client/program/ClientTest.java
index 9a37dde..594525e 100644
--- a/flink-clients/src/test/java/org/apache/flink/client/program/ClientTest.java
+++ b/flink-clients/src/test/java/org/apache/flink/client/program/ClientTest.java
@@ -21,7 +21,6 @@ package org.apache.flink.client.program;
 import akka.actor.ActorSystem;
 import akka.actor.Props;
 import akka.actor.Status;
-import akka.actor.UntypedActor;
 import org.apache.flink.api.common.InvalidProgramException;
 import org.apache.flink.api.common.JobSubmissionResult;
 import org.apache.flink.api.common.Plan;
@@ -34,6 +33,7 @@ import org.apache.flink.optimizer.plantranslate.JobGraphGenerator;
 import org.apache.flink.configuration.ConfigConstants;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.akka.AkkaUtils;
+import org.apache.flink.runtime.akka.FlinkUntypedActor;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.runtime.jobmanager.JobManager;
@@ -48,9 +48,12 @@ import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
+import scala.Option;
 import scala.Some;
 import scala.Tuple2;
 
+import java.util.UUID;
+
 import static org.junit.Assert.*;
 import static org.mockito.Matchers.any;
 
@@ -220,25 +223,49 @@ public class ClientTest {
 
 	// --------------------------------------------------------------------------------------------
 
-	public static class SuccessReturningActor extends UntypedActor {
+	public static class SuccessReturningActor extends FlinkUntypedActor {
+
+		private Option<UUID> leaderSessionID = Option.<UUID>apply(UUID.randomUUID());
 
 		@Override
-		public void onReceive(Object message) throws Exception {
+		public void handleMessage(Object message) {
 			if (message instanceof JobManagerMessages.SubmitJob) {
 				JobID jid = ((JobManagerMessages.SubmitJob) message).jobGraph().getJobID();
-				getSender().tell(new Status.Success(jid), getSelf());
+				getSender().tell(
+						decorateMessage(new Status.Success(jid)),
+						getSelf());
+			} else if(message instanceof JobManagerMessages.RequestLeaderSessionID$) {
+				getSender().tell(
+						decorateMessage(new JobManagerMessages.ResponseLeaderSessionID(leaderSessionID)),
+						getSelf());
 			}
 			else {
-				getSender().tell(new Status.Failure(new Exception("Unknown message " + message)), getSelf());
+				getSender().tell(
+						decorateMessage(new Status.Failure(new Exception("Unknown message " + message))),
+						getSelf());
 			}
 		}
+
+		@Override
+		protected Option<UUID> getLeaderSessionID() {
+			return leaderSessionID;
+		}
 	}
 
-	public static class FailureReturningActor extends UntypedActor {
+	public static class FailureReturningActor extends FlinkUntypedActor {
+
+		private Option<UUID> leaderSessionID = Option.<UUID>apply(UUID.randomUUID());
+
+		@Override
+		public void handleMessage(Object message) {
+			getSender().tell(
+					decorateMessage(new Status.Failure(new Exception("test"))),
+					getSelf());
+		}
 
 		@Override
-		public void onReceive(Object message) throws Exception {
-			getSender().tell(new Status.Failure(new Exception("test")), getSelf());
+		protected Option<UUID> getLeaderSessionID() {
+			return leaderSessionID;
 		}
 	}
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/akka/FlinkUntypedActor.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/akka/FlinkUntypedActor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/akka/FlinkUntypedActor.java
new file mode 100644
index 0000000..bba2aeb
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/akka/FlinkUntypedActor.java
@@ -0,0 +1,140 @@
+/*
+ * 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.akka;
+
+import akka.actor.UntypedActor;
+import org.apache.flink.runtime.messages.JobManagerMessages.LeaderSessionMessage;
+import org.apache.flink.runtime.messages.RequiresLeaderSessionID;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import scala.Option;
+
+import java.util.UUID;
+
+/**
+ * Base class for Flink's actors implemented with Java. Actors inheriting from this class
+ * automatically log received messages when the debug log level is activated. Furthermore,
+ * they filter out {@link LeaderSessionMessage} with the wrong leader session ID. If a message
+ * of type {@link RequiresLeaderSessionID} without being wrapped in a LeaderSessionMessage is
+ * detected, then an Exception is thrown.
+ *
+ * In order to implement the actor behavior, an implementing subclass has to override the method
+ * handleMessage, which defines how messages are processed. Furthermore, the subclass has to provide
+ * a leader session ID option which is returned by getLeaderSessionID.
+ */
+public abstract class FlinkUntypedActor extends UntypedActor {
+	protected static Logger LOG = LoggerFactory.getLogger(FlinkUntypedActor.class);
+
+	/**
+	 * This method is called by Akka if a new message has arrived for the actor. It logs the
+	 * processing time of the incoming message if the logging level is set to debug. After logging
+	 * the handleLeaderSessionID method is called.
+	 *
+	 * Important: This method cannot be overriden. The actor specific message handling logic is
+	 * implemented by the method handleMessage.
+	 *
+	 * @param message Incoming message
+	 * @throws Exception
+	 */
+	@Override
+	public final void onReceive(Object message) throws Exception {
+		if(LOG.isDebugEnabled()) {
+			LOG.debug("Received message {} at {} from {}.", message, getSelf().path(), getSender());
+
+			long start = System.nanoTime();
+
+			handleLeaderSessionID(message);
+
+			long duration = (System.nanoTime() - start)/ 1000000;
+
+			LOG.debug("Handled message {} in {} ms from {}.", message, duration, getSender());
+		} else {
+			handleLeaderSessionID(message);
+		}
+	}
+
+	/**
+	 * This method filters out {@link LeaderSessionMessage} whose leader session ID is not equal
+	 * to the actors leader session ID. If a message of type {@link RequiresLeaderSessionID}
+	 * arrives, then an Exception is thrown, because these messages have to be wrapped in a
+	 * {@link LeaderSessionMessage}.
+	 *
+	 * @param message Incoming message
+	 * @throws Exception
+	 */
+	private void handleLeaderSessionID(Object message) throws Exception {
+		if(message instanceof LeaderSessionMessage) {
+			LeaderSessionMessage msg = (LeaderSessionMessage) message;
+
+			if(msg.leaderSessionID().isDefined() && getLeaderSessionID().isDefined()) {
+				if(getLeaderSessionID().equals(msg.leaderSessionID())) {
+					// finally call method to handle message
+					handleMessage(msg.message());
+				} else {
+					handleDiscardedMessage(msg);
+				}
+			} else {
+				handleDiscardedMessage(msg);
+			}
+		} else if (message instanceof RequiresLeaderSessionID) {
+			throw new Exception("Received a message " + message + " without a leader session " +
+					"ID, even though it requires to have one.");
+		} else {
+			// call method to handle message
+			handleMessage(message);
+		}
+	}
+
+	private void handleDiscardedMessage(Object msg) {
+		LOG.debug("Discard message {} because the leader session ID was not correct.", msg);
+	}
+
+	/**
+	 * This method contains the actor logic which defines how to react to incoming messages.
+	 *
+	 * @param message Incoming message
+	 * @throws Exception
+	 */
+	protected abstract void handleMessage(Object message) throws Exception;
+
+	/**
+	 * Returns the current leader session ID associcated with this actor.
+	 * @return
+	 */
+	protected abstract Option<UUID> getLeaderSessionID();
+
+	/**
+	 * This method should be called for every outgoing message. It wraps messages which require
+	 * a leader session ID (indicated by {@link RequiresLeaderSessionID}) in a
+	 * {@link LeaderSessionMessage} with the actor's leader session ID.
+	 *
+	 * This method can be overriden to implement a different decoration behavior.
+	 *
+	 * @param message Message to be decorated
+	 * @return The deocrated message
+	 */
+	protected Object decorateMessage(Object message) {
+		if(message instanceof  RequiresLeaderSessionID) {
+			return new LeaderSessionMessage(getLeaderSessionID(), message);
+		} else {
+			return message;
+		}
+	}
+
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
index 2b2bf6b..9694132 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
@@ -18,7 +18,6 @@
 
 package org.apache.flink.runtime.checkpoint;
 
-import akka.actor.ActorRef;
 import akka.actor.ActorSystem;
 import akka.actor.PoisonPill;
 import akka.actor.Props;
@@ -28,12 +27,15 @@ import org.apache.flink.runtime.executiongraph.Execution;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
 import org.apache.flink.runtime.executiongraph.ExecutionJobVertex;
 import org.apache.flink.runtime.executiongraph.ExecutionVertex;
+import org.apache.flink.runtime.instance.AkkaActorGateway;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
 import org.apache.flink.runtime.messages.checkpoint.AcknowledgeCheckpoint;
 import org.apache.flink.runtime.messages.checkpoint.NotifyCheckpointComplete;
 import org.apache.flink.runtime.messages.checkpoint.TriggerCheckpoint;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import scala.Option;
 
 import java.util.ArrayDeque;
 import java.util.ArrayList;
@@ -44,6 +46,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Timer;
 import java.util.TimerTask;
+import java.util.UUID;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 
@@ -94,7 +97,7 @@ public class CheckpointCoordinator {
 	
 	private TimerTask periodicScheduler;
 	
-	private ActorRef jobStatusListener;
+	private ActorGateway jobStatusListener;
 	
 	private ClassLoader userClassLoader;
 	
@@ -102,11 +105,14 @@ public class CheckpointCoordinator {
 	
 	// --------------------------------------------------------------------------------------------
 
-	public CheckpointCoordinator(JobID job, int numSuccessfulCheckpointsToRetain, long checkpointTimeout,
-								ExecutionVertex[] tasksToTrigger,
-								ExecutionVertex[] tasksToWaitFor,
-								ExecutionVertex[] tasksToCommitTo,
-								ClassLoader userClassLoader) {
+	public CheckpointCoordinator(
+			JobID job,
+			int numSuccessfulCheckpointsToRetain,
+			long checkpointTimeout,
+			ExecutionVertex[] tasksToTrigger,
+			ExecutionVertex[] tasksToWaitFor,
+			ExecutionVertex[] tasksToCommitTo,
+			ClassLoader userClassLoader) {
 		
 		// some sanity checks
 		if (job == null || tasksToTrigger == null ||
@@ -157,7 +163,7 @@ public class CheckpointCoordinator {
 			
 			// make sure that the actor does not linger
 			if (jobStatusListener != null) {
-				jobStatusListener.tell(PoisonPill.getInstance(), ActorRef.noSender());
+				jobStatusListener.tell(PoisonPill.getInstance());
 				jobStatusListener = null;
 			}
 			
@@ -521,16 +527,28 @@ public class CheckpointCoordinator {
 		}
 	}
 	
-	public ActorRef createJobStatusListener(ActorSystem actorSystem, long checkpointInterval) {
+	public ActorGateway createJobStatusListener(
+			ActorSystem actorSystem,
+			long checkpointInterval,
+			Option<UUID> leaderSessionID) {
 		synchronized (lock) {
 			if (shutdown) {
 				throw new IllegalArgumentException("Checkpoint coordinator is shut down");
 			}
 
 			if (jobStatusListener == null) {
-				Props props = Props.create(CheckpointCoordinatorDeActivator.class, this, checkpointInterval);
-				jobStatusListener = actorSystem.actorOf(props);
+				Props props = Props.create(
+						CheckpointCoordinatorDeActivator.class,
+						this,
+						checkpointInterval,
+						leaderSessionID);
+
+				// wrap the ActorRef in a AkkaActorGateway to support message decoration
+				jobStatusListener = new AkkaActorGateway(
+						actorSystem.actorOf(props),
+						leaderSessionID);
 			}
+
 			return jobStatusListener;
 		}
 	}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorDeActivator.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorDeActivator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorDeActivator.java
index a6c4d76..f65be15 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorDeActivator.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorDeActivator.java
@@ -18,26 +18,38 @@
 
 package org.apache.flink.runtime.checkpoint;
 
-import akka.actor.UntypedActor;
+import com.google.common.base.Preconditions;
+import org.apache.flink.runtime.akka.FlinkUntypedActor;
 import org.apache.flink.runtime.jobgraph.JobStatus;
 import org.apache.flink.runtime.messages.ExecutionGraphMessages;
+import scala.Option;
+
+import java.util.UUID;
 
 /**
  * This actor listens to changes in the JobStatus and activates or deactivates the periodic
  * checkpoint scheduler.
  */
-public class CheckpointCoordinatorDeActivator extends UntypedActor {
+public class CheckpointCoordinatorDeActivator extends FlinkUntypedActor {
 
 	private final CheckpointCoordinator coordinator;
 	private final long interval;
+	private final Option<UUID> leaderSessionID;
 	
-	public CheckpointCoordinatorDeActivator(CheckpointCoordinator coordinator, long interval) {
+	public CheckpointCoordinatorDeActivator(
+			CheckpointCoordinator coordinator,
+			long interval,
+			Option<UUID> leaderSessionID) {
+		Preconditions.checkNotNull(coordinator, "The checkpointCoordinator must not be null.");
+		Preconditions.checkNotNull(leaderSessionID, "The leaderSesssionID must not be null.");
+
 		this.coordinator = coordinator;
 		this.interval = interval;
+		this.leaderSessionID = leaderSessionID;
 	}
 
 	@Override
-	public void onReceive(Object message) {
+	public void handleMessage(Object message) {
 		if (message instanceof ExecutionGraphMessages.JobStatusChanged) {
 			JobStatus status = ((ExecutionGraphMessages.JobStatusChanged) message).newJobStatus();
 			
@@ -53,4 +65,9 @@ public class CheckpointCoordinatorDeActivator extends UntypedActor {
 		
 		// we ignore all other messages
 	}
+
+	@Override
+	public Option<UUID> getLeaderSessionID() {
+		return leaderSessionID;
+	}
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/client/JobClient.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/client/JobClient.java b/flink-runtime/src/main/java/org/apache/flink/runtime/client/JobClient.java
index 72f39d6..aeefa61 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/client/JobClient.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/client/JobClient.java
@@ -26,10 +26,12 @@ import akka.actor.Props;
 import akka.actor.Status;
 import akka.pattern.Patterns;
 import akka.util.Timeout;
+import com.google.common.base.Preconditions;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.configuration.ConfigConstants;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.akka.AkkaUtils;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.runtime.messages.JobClientMessages;
 import org.apache.flink.runtime.messages.JobManagerMessages;
@@ -114,25 +116,36 @@ public class JobClient {
 	 * case a [[JobExecutionException]] is thrown.
 	 *
 	 * @param actorSystem The actor system that performs the communication.
-	 * @param jobManager  The JobManager that should execute the job.
+	 * @param jobManagerGateway  Gateway to the JobManager that should execute the job.
 	 * @param jobGraph    JobGraph describing the Flink job
 	 * @param timeout     Timeout for futures
 	 * @return The job execution result
 	 * @throws org.apache.flink.runtime.client.JobExecutionException Thrown if the job
 	 *                                                               execution fails.
 	 */
-	public static SerializedJobExecutionResult submitJobAndWait(ActorSystem actorSystem, ActorRef jobManager,
-																JobGraph jobGraph, FiniteDuration timeout,
-																boolean sysoutLogUpdates) throws JobExecutionException
-	{
-		if (actorSystem == null || jobManager == null || jobGraph == null || timeout == null) {
-			throw new NullPointerException();
-		}
+	public static SerializedJobExecutionResult submitJobAndWait(
+			ActorSystem actorSystem,
+			ActorGateway jobManagerGateway,
+			JobGraph jobGraph,
+			FiniteDuration timeout,
+			boolean sysoutLogUpdates) throws JobExecutionException {
+
+		Preconditions.checkNotNull(actorSystem, "The actorSystem must not be null.");
+		Preconditions.checkNotNull(jobManagerGateway, "The jobManagerGateway must not be null.");
+		Preconditions.checkNotNull(jobGraph, "The jobGraph must not be null.");
+		Preconditions.checkNotNull(timeout, "The timeout must not be null.");
+
 		// for this job, we create a proxy JobClientActor that deals with all communication with
 		// the JobManager. It forwards the job submission, checks the success/failure responses, logs
 		// update messages, watches for disconnect between client and JobManager, ...
 
-		Props jobClientActorProps = Props.create(JobClientActor.class, jobManager, LOG, sysoutLogUpdates);
+		Props jobClientActorProps = Props.create(
+				JobClientActor.class,
+				jobManagerGateway.actor(),
+				LOG,
+				sysoutLogUpdates,
+				jobManagerGateway.leaderSessionID());
+
 		ActorRef jobClientActor = actorSystem.actorOf(jobClientActorProps);
 
 		try {
@@ -178,25 +191,30 @@ public class JobClient {
 	 * Submits a job in detached mode. The method sends the JobGraph to the
 	 * JobManager and waits for the answer whether teh job could be started or not.
 	 *
+	 * @param jobManagerGateway Gateway to the JobManager which will execute the jobs
 	 * @param jobGraph The job
 	 * @param timeout  Timeout in which the JobManager must have responded.
 	 */
-	public static void submitJobDetached(ActorRef jobManager, JobGraph jobGraph, FiniteDuration timeout)
-			throws JobExecutionException {
-		if (jobManager == null || jobGraph == null || timeout == null) {
-			throw new NullPointerException();
-		}
+	public static void submitJobDetached(
+			ActorGateway jobManagerGateway,
+			JobGraph jobGraph,
+			FiniteDuration timeout) throws JobExecutionException {
+
+		Preconditions.checkNotNull(jobManagerGateway, "The jobManagerGateway must not be null.");
+		Preconditions.checkNotNull(jobGraph, "The jobGraph must not be null.");
+		Preconditions.checkNotNull(timeout, "The timeout must not be null.");
 
-		Future<Object> future = Patterns.ask(jobManager,
+		Future<Object> future = jobManagerGateway.ask(
 				new JobManagerMessages.SubmitJob(jobGraph, false),
-				new Timeout(timeout));
+				timeout);
+
 		try {
 			Object result = Await.result(future, timeout);
 			if (result instanceof JobID) {
 				JobID respondedID = (JobID) result;
 				if (!respondedID.equals(jobGraph.getJobID())) {
-					throw new Exception("JobManager responded for wrong Job. This Job: " + jobGraph.getJobID() +
-							", response: " + respondedID);
+					throw new Exception("JobManager responded for wrong Job. This Job: " +
+							jobGraph.getJobID() + ", response: " + respondedID);
 				}
 			}
 			else {
@@ -222,17 +240,20 @@ public class JobClient {
 	 * blocking call.
 	 *
 	 * @param jobGraph   Flink job containing the information about the required jars
-	 * @param jobManager ActorRef of the JobManager.
+	 * @param jobManagerGateway Gateway to the JobManager.
 	 * @param timeout    Timeout for futures
 	 * @throws IOException Thrown, if the file upload to the JobManager failed.
 	 */
-	public static void uploadJarFiles(JobGraph jobGraph, ActorRef jobManager, FiniteDuration timeout)
+	public static void uploadJarFiles(
+			JobGraph jobGraph,
+			ActorGateway jobManagerGateway,
+			FiniteDuration timeout)
 			throws IOException {
 		if (jobGraph.hasUsercodeJarFiles()) {
-			Timeout tOut = new Timeout(timeout);
-			Future<Object> futureBlobPort = Patterns.ask(jobManager,
+
+			Future<Object> futureBlobPort = jobManagerGateway.ask(
 					JobManagerMessages.getRequestBlobManagerPort(),
-					tOut);
+					timeout);
 
 			int port;
 			try {
@@ -247,7 +268,7 @@ public class JobClient {
 				throw new IOException("Could not retrieve the JobManager's blob port.", e);
 			}
 
-			Option<String> jmHost = jobManager.path().address().host();
+			Option<String> jmHost = jobManagerGateway.actor().path().address().host();
 			String jmHostname = jmHost.isDefined() ? jmHost.get() : "localhost";
 			InetSocketAddress serverAddress = new InetSocketAddress(jmHostname, port);
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/client/JobClientActor.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/client/JobClientActor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/client/JobClientActor.java
index ee31e8d..b4bfc8f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/client/JobClientActor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/client/JobClientActor.java
@@ -22,37 +22,47 @@ import akka.actor.ActorRef;
 import akka.actor.PoisonPill;
 import akka.actor.Status;
 import akka.actor.Terminated;
-import akka.actor.UntypedActor;
+import com.google.common.base.Preconditions;
+import org.apache.flink.runtime.akka.FlinkUntypedActor;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.runtime.messages.ExecutionGraphMessages;
 import org.apache.flink.runtime.messages.JobClientMessages;
 import org.apache.flink.runtime.messages.JobManagerMessages;
 import org.slf4j.Logger;
+import scala.Option;
+
+import java.util.UUID;
 
 /**
  * Actor which constitutes the bridge between the non-actor code and the JobManager. The JobClient
  * is used to submit jobs to the JobManager and to request the port of the BlobManager.
  */
-public class JobClientActor extends UntypedActor {
+public class JobClientActor extends FlinkUntypedActor {
 	
 	private final ActorRef jobManager;
 	private final Logger logger;
 	private final boolean sysoutUpdates;
-	
+
+	// leader session ID of the JobManager when this actor was created
+	private final Option<UUID> leaderSessionID;
+
+	// Actor which submits a job to the JobManager via this actor
 	private ActorRef submitter;
-	
-	
-	public JobClientActor(ActorRef jobManager, Logger logger, boolean sysoutUpdates) {
-		if (jobManager == null || logger == null) {
-			throw new NullPointerException();
-		}
-		this.jobManager = jobManager;
-		this.logger = logger;
+
+	public JobClientActor(
+			ActorRef jobManager,
+			Logger logger,
+			boolean sysoutUpdates,
+			Option<UUID> leaderSessionID) {
+		this.jobManager = Preconditions.checkNotNull(jobManager, "The JobManager ActorRef must not be null.");
+		this.logger = Preconditions.checkNotNull(logger, "The logger must not be null.");
+		this.leaderSessionID = Preconditions.checkNotNull(leaderSessionID, "The leader session ID option must not be null.");
+
 		this.sysoutUpdates = sysoutUpdates;
 	}
 	
 	@Override
-	public void onReceive(Object message) {
+	protected void handleMessage(Object message) {
 		
 		// =========== State Change Messages ===============
 
@@ -73,14 +83,17 @@ public class JobClientActor extends UntypedActor {
 				JobGraph jobGraph = ((JobClientMessages.SubmitJobAndWait) message).jobGraph();
 				if (jobGraph == null) {
 					logger.error("Received null JobGraph");
-					sender().tell(new Status.Failure(new Exception("JobGraph is null")), getSelf());
+					sender().tell(
+							decorateMessage(new Status.Failure(new Exception("JobGraph is null"))),
+							getSelf());
 				}
 				else {
 					logger.info("Sending message to JobManager {} to submit job {} ({}) and wait for progress",
 							jobManager.path().toString(), jobGraph.getName(), jobGraph.getJobID());
 
 					this.submitter = getSender();
-					jobManager.tell(new JobManagerMessages.SubmitJob(jobGraph, true), getSelf());
+					jobManager.tell(
+							decorateMessage(new JobManagerMessages.SubmitJob(jobGraph, true)), getSelf());
 					
 					// make sure we notify the sender when the connection got lost
 					getContext().watch(jobManager);
@@ -90,10 +103,12 @@ public class JobClientActor extends UntypedActor {
 				// repeated submission - tell failure to sender and kill self
 				String msg = "Received repeated 'SubmitJobAndWait'";
 				logger.error(msg);
-				getSender().tell(new Status.Failure(new Exception(msg)), ActorRef.noSender());
+				getSender().tell(
+						decorateMessage(new Status.Failure(new Exception(msg))),
+						ActorRef.noSender());
 
 				getContext().unwatch(jobManager);
-				getSelf().tell(PoisonPill.getInstance(), ActorRef.noSender());
+				getSelf().tell(decorateMessage(PoisonPill.getInstance()), ActorRef.noSender());
 			}
 		}
 		// acknowledgement to submit job is only logged, our original
@@ -102,12 +117,12 @@ public class JobClientActor extends UntypedActor {
 			// forward the success to the original job submitter
 			logger.debug("Received JobResultSuccess message from JobManager");
 			if (this.submitter != null) {
-				this.submitter.tell(message, getSelf());
+				this.submitter.tell(decorateMessage(message), getSelf());
 			}
 			
 			// we are done, stop ourselves
 			getContext().unwatch(jobManager);
-			getSelf().tell(PoisonPill.getInstance(), ActorRef.noSender());
+			getSelf().tell(decorateMessage(PoisonPill.getInstance()), ActorRef.noSender());
 		}
 		else if (message instanceof Status.Success) {
 			// job was successfully submitted :-)
@@ -117,7 +132,7 @@ public class JobClientActor extends UntypedActor {
 			// job execution failed, inform the actor that submitted the job
 			logger.debug("Received failure from JobManager", ((Status.Failure) message).cause());
 			if (submitter != null) {
-				submitter.tell(message, sender());
+				submitter.tell(decorateMessage(message), sender());
 			}
 		}
 
@@ -128,7 +143,7 @@ public class JobClientActor extends UntypedActor {
 			if (jobManager.equals(target)) {
 				String msg = "Lost connection to JobManager " + jobManager.path();
 				logger.info(msg);
-				submitter.tell(new Status.Failure(new Exception(msg)), getSelf());
+				submitter.tell(decorateMessage(new Status.Failure(new Exception(msg))), getSelf());
 			} else {
 				logger.error("Received 'Terminated' for unknown actor " + target);
 			}
@@ -140,7 +155,12 @@ public class JobClientActor extends UntypedActor {
 			logger.error("JobClient received unknown message: " + message);
 		}
 	}
-	
+
+	@Override
+	protected Option<UUID> getLeaderSessionID() {
+		return leaderSessionID;
+	}
+
 	private void logAndPrintMessage(Object message) {
 		logger.info(message.toString());
 		if (sysoutUpdates) {

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
index 3b836ad..aa0f981 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
@@ -30,7 +30,7 @@ import org.apache.flink.runtime.deployment.TaskDeploymentDescriptor;
 import org.apache.flink.runtime.execution.ExecutionState;
 import org.apache.flink.runtime.instance.Instance;
 import org.apache.flink.runtime.instance.InstanceConnectionInfo;
-import org.apache.flink.runtime.instance.InstanceGateway;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.instance.SimpleSlot;
 import org.apache.flink.runtime.io.network.ConnectionID;
 import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
@@ -361,7 +361,7 @@ public class Execution implements Serializable {
 			vertex.getExecutionGraph().registerExecution(this);
 
 			final Instance instance = slot.getInstance();
-			final InstanceGateway gateway = instance.getInstanceGateway();
+			final ActorGateway gateway = instance.getActorGateway();
 
 			final Future<Object> deployAction = gateway.ask(new SubmitTask(deployment), timeout);
 
@@ -848,7 +848,7 @@ public class Execution implements Serializable {
 
 		if (slot != null) {
 
-			final InstanceGateway gateway = slot.getInstance().getInstanceGateway();
+			final ActorGateway gateway = slot.getInstance().getActorGateway();
 
 			Future<Object> cancelResult = gateway.retry(
 				new CancelTask(attemptId),
@@ -881,7 +881,7 @@ public class Execution implements Serializable {
 			final Instance instance = slot.getInstance();
 
 			if (instance.isAlive()) {
-				final InstanceGateway gateway = instance.getInstanceGateway();
+				final ActorGateway gateway = instance.getActorGateway();
 
 				// TODO For some tests this could be a problem when querying too early if all resources were released
 				gateway.tell(new FailIntermediateResultPartitions(attemptId));
@@ -901,7 +901,7 @@ public class Execution implements Serializable {
 
 		if (consumerSlot != null) {
 			final Instance instance = consumerSlot.getInstance();
-			final InstanceGateway gateway = instance.getInstanceGateway();
+			final ActorGateway gateway = instance.getActorGateway();
 
 			Future<Object> futureUpdate = gateway.ask(updatePartitionInfo, timeout);
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java
index 9c977fa..833518c 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java
@@ -18,8 +18,6 @@
 
 package org.apache.flink.runtime.executiongraph;
 
-import akka.actor.ActorRef;
-
 import akka.actor.ActorSystem;
 import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.accumulators.Accumulator;
@@ -32,6 +30,7 @@ import org.apache.flink.runtime.accumulators.StringifiedAccumulatorResult;
 import org.apache.flink.runtime.blob.BlobKey;
 import org.apache.flink.runtime.checkpoint.CheckpointCoordinator;
 import org.apache.flink.runtime.execution.ExecutionState;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
 import org.apache.flink.runtime.jobgraph.JobVertex;
 import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
@@ -50,6 +49,7 @@ import org.apache.flink.util.InstantiationUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import scala.Option;
 import scala.concurrent.ExecutionContext;
 import scala.concurrent.duration.FiniteDuration;
 
@@ -63,6 +63,7 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.UUID;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
@@ -166,10 +167,10 @@ public class ExecutionGraph implements Serializable {
 
 	/** Listeners that receive messages when the entire job switches it status (such as from
 	 * RUNNING to FINISHED) */
-	private final List<ActorRef> jobStatusListenerActors;
+	private final List<ActorGateway> jobStatusListenerActors;
 
 	/** Listeners that receive messages whenever a single task execution changes its status */
-	private final List<ActorRef> executionListenerActors;
+	private final List<ActorGateway> executionListenerActors;
 
 	/** Timestamps (in milliseconds as returned by {@code System.currentTimeMillis()} when
 	 * the execution graph transitioned into a certain state. The index into this array is the
@@ -287,8 +288,8 @@ public class ExecutionGraph implements Serializable {
 		this.verticesInCreationOrder = new ArrayList<ExecutionJobVertex>();
 		this.currentExecutions = new ConcurrentHashMap<ExecutionAttemptID, Execution>();
 
-		this.jobStatusListenerActors  = new CopyOnWriteArrayList<ActorRef>();
-		this.executionListenerActors = new CopyOnWriteArrayList<ActorRef>();
+		this.jobStatusListenerActors  = new CopyOnWriteArrayList<ActorGateway>();
+		this.executionListenerActors = new CopyOnWriteArrayList<ActorGateway>();
 
 		this.stateTimestamps = new long[JobStatus.values().length];
 		this.stateTimestamps[JobStatus.CREATED.ordinal()] = System.currentTimeMillis();
@@ -340,12 +341,14 @@ public class ExecutionGraph implements Serializable {
 		return scheduleMode;
 	}
 
-	public void enableSnaphotCheckpointing(long interval, long checkpointTimeout,
-											List<ExecutionJobVertex> verticesToTrigger,
-											List<ExecutionJobVertex> verticesToWaitFor,
-											List<ExecutionJobVertex> verticesToCommitTo,
-											ActorSystem actorSystem)
-	{
+	public void enableSnapshotCheckpointing(
+			long interval,
+			long checkpointTimeout,
+			List<ExecutionJobVertex> verticesToTrigger,
+			List<ExecutionJobVertex> verticesToWaitFor,
+			List<ExecutionJobVertex> verticesToCommitTo,
+			ActorSystem actorSystem,
+			Option<UUID> leaderSessionID) {
 		// simple sanity checks
 		if (interval < 10 || checkpointTimeout < 10) {
 			throw new IllegalArgumentException();
@@ -363,12 +366,22 @@ public class ExecutionGraph implements Serializable {
 		
 		// create the coordinator that triggers and commits checkpoints and holds the state 
 		snapshotCheckpointsEnabled = true;
-		checkpointCoordinator = new CheckpointCoordinator(jobID, NUMBER_OF_SUCCESSFUL_CHECKPOINTS_TO_RETAIN,
-				checkpointTimeout, tasksToTrigger, tasksToWaitFor, tasksToCommitTo, userClassLoader);
+		checkpointCoordinator = new CheckpointCoordinator(
+				jobID,
+				NUMBER_OF_SUCCESSFUL_CHECKPOINTS_TO_RETAIN,
+				checkpointTimeout,
+				tasksToTrigger,
+				tasksToWaitFor,
+				tasksToCommitTo,
+				userClassLoader);
 		
 		// the periodic checkpoint scheduler is activated and deactivated as a result of
 		// job status changes (running -> on, all other states -> off)
-		registerJobStatusListener(checkpointCoordinator.createJobStatusListener(actorSystem, interval));
+		registerJobStatusListener(
+				checkpointCoordinator.createJobStatusListener(
+						actorSystem,
+						interval,
+						leaderSessionID));
 	}
 	
 	public void disableSnaphotCheckpointing() {
@@ -998,13 +1011,13 @@ public class ExecutionGraph implements Serializable {
 	//  Listeners & Observers
 	// --------------------------------------------------------------------------------------------
 
-	public void registerJobStatusListener(ActorRef listener) {
+	public void registerJobStatusListener(ActorGateway listener) {
 		if (listener != null) {
 			this.jobStatusListenerActors.add(listener);
 		}
 	}
 
-	public void registerExecutionListener(ActorRef listener) {
+	public void registerExecutionListener(ActorGateway listener) {
 		if (listener != null) {
 			this.executionListenerActors.add(listener);
 		}
@@ -1016,8 +1029,8 @@ public class ExecutionGraph implements Serializable {
 			ExecutionGraphMessages.JobStatusChanged message =
 					new ExecutionGraphMessages.JobStatusChanged(jobID, newState, System.currentTimeMillis(), error);
 
-			for (ActorRef listener: jobStatusListenerActors) {
-				listener.tell(message, ActorRef.noSender());
+			for (ActorGateway listener: jobStatusListenerActors) {
+				listener.tell(message);
 			}
 		}
 	}
@@ -1035,8 +1048,8 @@ public class ExecutionGraph implements Serializable {
 																	executionID, newExecutionState,
 																	System.currentTimeMillis(), message);
 
-			for (ActorRef listener : executionListenerActors) {
-				listener.tell(actorMessage, ActorRef.noSender());
+			for (ActorGateway listener : executionListenerActors) {
+				listener.tell(actorMessage);
 			}
 		}
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
index f9001cf..dcf64c0 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
@@ -28,7 +28,7 @@ import org.apache.flink.runtime.deployment.TaskDeploymentDescriptor;
 import org.apache.flink.runtime.execution.ExecutionState;
 import org.apache.flink.runtime.instance.Instance;
 import org.apache.flink.runtime.instance.InstanceConnectionInfo;
-import org.apache.flink.runtime.instance.InstanceGateway;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.instance.SimpleSlot;
 import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
 import org.apache.flink.runtime.jobgraph.DistributionPattern;
@@ -469,7 +469,7 @@ public class ExecutionVertex implements Serializable {
 			
 			// send only if we actually have a target
 			if (slot != null) {
-				InstanceGateway gateway = slot.getInstance().getInstanceGateway();
+				ActorGateway gateway = slot.getInstance().getActorGateway();
 				if (gateway != null) {
 					gateway.tell(message);
 				}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/instance/ActorGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/instance/ActorGateway.java b/flink-runtime/src/main/java/org/apache/flink/runtime/instance/ActorGateway.java
new file mode 100644
index 0000000..fe4a1cd
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/instance/ActorGateway.java
@@ -0,0 +1,107 @@
+/*
+ * 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.instance;
+
+import akka.actor.ActorRef;
+import scala.Option;
+import scala.concurrent.ExecutionContext;
+import scala.concurrent.Future;
+import scala.concurrent.duration.FiniteDuration;
+
+import java.util.UUID;
+
+/**
+ * Interface to abstract the communication with an actor.
+ *
+ * It allows to avoid direct interaction with an ActorRef.
+ */
+public interface ActorGateway {
+
+	/**
+	 * Sends a message asynchronously and returns its response. The response to the message is
+	 * returned as a future.
+	 *
+	 * @param message Message to be sent
+	 * @param timeout Timeout until the Future is completed with an AskTimeoutException
+	 * @return Future which contains the response to the sent message
+	 */
+	Future<Object> ask(Object message, FiniteDuration timeout);
+
+	/**
+	 * Sends a message asynchronously without a result.
+	 *
+	 * @param message Message to be sent
+	 */
+	void tell(Object message);
+
+	/**
+	 * Sends a message asynchronously without a result with sender being the sender.
+	 *
+	 * @param message Message to be sent
+	 * @param sender Sender of the message
+	 */
+	void tell(Object message, ActorGateway sender);
+
+	/**
+	 * Forwards a message. For the receiver of this message it looks as if sender has sent the
+	 * message.
+	 *
+	 * @param message Message to be sent
+	 * @param sender Sender of the forwarded message
+	 */
+	void forward(Object message, ActorGateway sender);
+
+	/**
+	 * Retries to send asynchronously a message up to numberRetries times. The response to this
+	 * message is returned as a future. The message is re-sent if the number of retries is not yet
+	 * exceeded and if an exception occurred while sending it.
+	 *
+	 * @param message Message to be sent
+	 * @param numberRetries Number of times to retry sending the message
+	 * @param timeout Timeout for each sending attempt
+	 * @param executionContext ExecutionContext which is used to send the message multiple times
+	 * @return Future of the response to the sent message
+	 */
+	Future<Object> retry(
+			Object message,
+			int numberRetries,
+			FiniteDuration timeout,
+			ExecutionContext executionContext);
+
+	/**
+	 * Returns the path of the remote instance.
+	 *
+	 * @return Path of the remote instance.
+	 */
+	String path();
+
+	/**
+	 * Returns the underlying actor with which is communicated
+	 *
+	 * @return ActorRef of the target actor
+	 */
+	ActorRef actor();
+
+	/**
+	 * Returns the leaderSessionID associated with the remote actor or None.
+	 *
+	 * @return Leader session ID if its associated with this gateway, otherwise None
+	 */
+	Option<UUID> leaderSessionID();
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/instance/AkkaActorGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/instance/AkkaActorGateway.java b/flink-runtime/src/main/java/org/apache/flink/runtime/instance/AkkaActorGateway.java
new file mode 100644
index 0000000..ea55458
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/instance/AkkaActorGateway.java
@@ -0,0 +1,157 @@
+/*
+ * 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.instance;
+
+import akka.actor.ActorRef;
+import akka.pattern.Patterns;
+import akka.util.Timeout;
+import org.apache.flink.runtime.LeaderSessionMessageDecorator;
+import org.apache.flink.runtime.MessageDecorator;
+import org.apache.flink.runtime.akka.AkkaUtils;
+import scala.Option;
+import scala.concurrent.ExecutionContext;
+import scala.concurrent.Future;
+import scala.concurrent.duration.FiniteDuration;
+
+import java.util.UUID;
+
+/**
+ * Concrete {@link ActorGateway} implementation which uses Akka to communicate with remote actors.
+ */
+public class AkkaActorGateway implements ActorGateway {
+
+	// ActorRef of the remote instance
+	private final ActorRef actor;
+
+	// Associated leader session ID, which is used for RequiresLeaderSessionID messages
+	private final Option<UUID> leaderSessionID;
+
+	// Decorator for messages
+	private final MessageDecorator decorator;
+
+	public AkkaActorGateway(ActorRef actor, Option<UUID> leaderSessionID) {
+		this.actor = actor;
+		this.leaderSessionID = leaderSessionID;
+		// we want to wrap RequiresLeaderSessionID messages in a LeaderSessionMessage
+		this.decorator = new LeaderSessionMessageDecorator(leaderSessionID);
+	}
+
+	/**
+	 * Sends a message asynchronously and returns its response. The response to the message is
+	 * returned as a future.
+	 *
+	 * @param message Message to be sent
+	 * @param timeout Timeout until the Future is completed with an AskTimeoutException
+	 * @return Future which contains the response to the sent message
+	 */
+	@Override
+	public Future<Object> ask(Object message, FiniteDuration timeout) {
+		Object newMessage = decorator.decorate(message);
+		return Patterns.ask(actor, newMessage, new Timeout(timeout));
+	}
+
+	/**
+	 * Sends a message asynchronously without a result.
+	 *
+	 * @param message Message to be sent
+	 */
+	@Override
+	public void tell(Object message) {
+		Object newMessage = decorator.decorate(message);
+		actor.tell(newMessage, ActorRef.noSender());
+	}
+
+	/**
+	 * Sends a message asynchronously without a result with sender being the sender.
+	 *
+	 * @param message Message to be sent
+	 * @param sender Sender of the message
+	 */
+	@Override
+	public void tell(Object message, ActorGateway sender) {
+		Object newMessage = decorator.decorate(message);
+		actor.tell(newMessage, sender.actor());
+	}
+
+	/**
+	 * Forwards a message. For the receiver of this message it looks as if sender has sent the
+	 * message.
+	 *
+	 * @param message Message to be sent
+	 * @param sender Sender of the forwarded message
+	 */
+	@Override
+	public void forward(Object message, ActorGateway sender) {
+		Object newMessage = decorator.decorate(message);
+		actor.tell(newMessage, sender.actor());
+	}
+
+	/**
+	 * Retries to send asynchronously a message up to numberRetries times. The response to this
+	 * message is returned as a future. The message is re-sent if the number of retries is not yet
+	 * exceeded and if an exception occurred while sending it.
+	 *
+	 * @param message Message to be sent
+	 * @param numberRetries Number of times to retry sending the message
+	 * @param timeout Timeout for each sending attempt
+	 * @param executionContext ExecutionContext which is used to send the message multiple times
+	 * @return Future of the response to the sent message
+	 */
+	@Override
+	public Future<Object> retry(
+			Object message,
+			int numberRetries,
+			FiniteDuration timeout,
+			ExecutionContext executionContext) {
+
+		Object newMessage = decorator.decorate(message);
+
+		return AkkaUtils.retry(
+			actor,
+			newMessage,
+			numberRetries,
+			executionContext,
+			timeout);
+	}
+
+	/**
+	 * Returns the ActorPath of the remote instance.
+	 *
+	 * @return ActorPath of the remote instance.
+	 */
+	@Override
+	public String path() {
+		return actor.path().toString();
+	}
+
+	/**
+	 * Returns {@link ActorRef} of the target actor
+	 *
+	 * @return ActorRef of the target actor
+	 */
+	@Override
+	public ActorRef actor() {
+		return actor;
+	}
+
+	@Override
+	public Option<UUID> leaderSessionID() {
+		return leaderSessionID;
+	}
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/instance/AkkaInstanceGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/instance/AkkaInstanceGateway.java b/flink-runtime/src/main/java/org/apache/flink/runtime/instance/AkkaInstanceGateway.java
deleted file mode 100644
index b7d60c5..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/instance/AkkaInstanceGateway.java
+++ /dev/null
@@ -1,111 +0,0 @@
-/*
- * 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.instance;
-
-import akka.actor.ActorRef;
-import akka.pattern.Patterns;
-import akka.util.Timeout;
-import org.apache.flink.runtime.akka.AkkaUtils;
-import scala.concurrent.ExecutionContext;
-import scala.concurrent.Future;
-import scala.concurrent.duration.FiniteDuration;
-
-/**
- * InstanceGateway implementation which uses Akka to communicate with remote instances.
- */
-public class AkkaInstanceGateway implements InstanceGateway {
-
-	/** ActorRef of the remote instance */
-	private final ActorRef taskManager;
-
-	public AkkaInstanceGateway(ActorRef taskManager) {
-		this.taskManager = taskManager;
-	}
-
-	/**
-	 * Sends a message asynchronously and returns its response. The response to the message is
-	 * returned as a future.
-	 *
-	 * @param message Message to be sent
-	 * @param timeout Timeout until the Future is completed with an AskTimeoutException
-	 * @return Future which contains the response to the sent message
-	 */
-	@Override
-	public Future<Object> ask(Object message, FiniteDuration timeout) {
-		return Patterns.ask(taskManager, message, new Timeout(timeout));
-	}
-
-	/**
-	 * Sends a message asynchronously without a result.
-	 *
-	 * @param message Message to be sent
-	 */
-	@Override
-	public void tell(Object message) {
-		taskManager.tell(message, ActorRef.noSender());
-	}
-
-	/**
-	 * Forwards a message. For the receiver of this message it looks as if sender has sent the
-	 * message.
-	 *
-	 * @param message Message to be sent
-	 * @param sender Sender of the forwarded message
-	 */
-	@Override
-	public void forward(Object message, ActorRef sender) {
-		taskManager.tell(message, sender);
-	}
-
-	/**
-	 * Retries to send asynchronously a message up to numberRetries times. The response to this
-	 * message is returned as a future. The message is re-sent if the number of retries is not yet
-	 * exceeded and if an exception occurred while sending it.
-	 *
-	 * @param message Message to be sent
-	 * @param numberRetries Number of times to retry sending the message
-	 * @param timeout Timeout for each sending attempt
-	 * @param executionContext ExecutionContext which is used to send the message multiple times
-	 * @return Future of the response to the sent message
-	 */
-	@Override
-	public Future<Object> retry(
-			Object message,
-			int numberRetries,
-			FiniteDuration timeout,
-			ExecutionContext executionContext) {
-
-		return AkkaUtils.retry(
-			taskManager,
-			message,
-			numberRetries,
-			executionContext,
-			timeout);
-	}
-
-	/**
-	 * Returns the ActorPath of the remote instance.
-	 *
-	 * @return ActorPath of the remote instance.
-	 */
-	@Override
-	public String path() {
-		return taskManager.path().toString();
-	}
-}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/instance/Instance.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/instance/Instance.java b/flink-runtime/src/main/java/org/apache/flink/runtime/instance/Instance.java
index 1c44b5d..743dc8f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/instance/Instance.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/instance/Instance.java
@@ -42,7 +42,7 @@ public class Instance {
 	private final Object instanceLock = new Object();
 
 	/** The instacne gateway to communicate with the instance */
-	private final InstanceGateway instanceGateway;
+	private final ActorGateway actorGateway;
 
 	/** The instance connection information for the data transfer. */
 	private final InstanceConnectionInfo connectionInfo;
@@ -79,19 +79,19 @@ public class Instance {
 	/**
 	 * Constructs an instance reflecting a registered TaskManager.
 	 *
-	 * @param instanceGateway The instance gateway to communicate with the remote instance
+	 * @param actorGateway The actor gateway to communicate with the remote instance
 	 * @param connectionInfo The remote connection where the task manager receives requests.
 	 * @param id The id under which the taskManager is registered.
 	 * @param resources The resources available on the machine.
 	 * @param numberOfSlots The number of task slots offered by this taskManager.
 	 */
 	public Instance(
-			InstanceGateway instanceGateway,
+			ActorGateway actorGateway,
 			InstanceConnectionInfo connectionInfo,
 			InstanceID id,
 			HardwareDescription resources,
 			int numberOfSlots) {
-		this.instanceGateway = instanceGateway;
+		this.actorGateway = actorGateway;
 		this.connectionInfo = connectionInfo;
 		this.instanceId = id;
 		this.resources = resources;
@@ -335,8 +335,8 @@ public class Instance {
 	 *
 	 * @return InstanceGateway associated with this instance
 	 */
-	public InstanceGateway getInstanceGateway() {
-		return instanceGateway;
+	public ActorGateway getActorGateway() {
+		return actorGateway;
 	}
 
 	public InstanceConnectionInfo getInstanceConnectionInfo() {
@@ -390,6 +390,6 @@ public class Instance {
 	@Override
 	public String toString() {
 		return String.format("%s @ %s - %d slots - URL: %s", instanceId, connectionInfo.getHostname(),
-				numberOfSlots, (instanceGateway != null ? instanceGateway.path() : "No instance gateway"));
+				numberOfSlots, (actorGateway != null ? actorGateway.path() : "No instance gateway"));
 	}
 }
\ No newline at end of file

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/instance/InstanceGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/instance/InstanceGateway.java b/flink-runtime/src/main/java/org/apache/flink/runtime/instance/InstanceGateway.java
deleted file mode 100644
index a30b2f6..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/instance/InstanceGateway.java
+++ /dev/null
@@ -1,82 +0,0 @@
-/*
- * 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.instance;
-
-import akka.actor.ActorRef;
-import scala.concurrent.ExecutionContext;
-import scala.concurrent.Future;
-import scala.concurrent.duration.FiniteDuration;
-
-/**
- * Interface to abstract the communication with an Instance.
- *
- * It allows to avoid direct interaction with an ActorRef.
- */
-public interface InstanceGateway {
-
-	/**
-	 * Sends a message asynchronously and returns its response. The response to the message is
-	 * returned as a future.
-	 *
-	 * @param message Message to be sent
-	 * @param timeout Timeout until the Future is completed with an AskTimeoutException
-	 * @return Future which contains the response to the sent message
-	 */
-	Future<Object> ask(Object message, FiniteDuration timeout);
-
-	/**
-	 * Sends a message asynchronously without a result.
-	 *
-	 * @param message Message to be sent
-	 */
-	void tell(Object message);
-
-	/**
-	 * Forwards a message. For the receiver of this message it looks as if sender has sent the
-	 * message.
-	 *
-	 * @param message Message to be sent
-	 * @param sender Sender of the forwarded message
-	 */
-	void forward(Object message, ActorRef sender);
-
-	/**
-	 * Retries to send asynchronously a message up to numberRetries times. The response to this
-	 * message is returned as a future. The message is re-sent if the number of retries is not yet
-	 * exceeded and if an exception occurred while sending it.
-	 *
-	 * @param message Message to be sent
-	 * @param numberRetries Number of times to retry sending the message
-	 * @param timeout Timeout for each sending attempt
-	 * @param executionContext ExecutionContext which is used to send the message multiple times
-	 * @return Future of the response to the sent message
-	 */
-	Future<Object> retry(
-			Object message,
-			int numberRetries,
-			FiniteDuration timeout,
-			ExecutionContext executionContext);
-
-	/**
-	 * Returns the path of the remote instance.
-	 *
-	 * @return Path of the remote instance.
-	 */
-	String path();
-}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/instance/InstanceManager.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/instance/InstanceManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/instance/InstanceManager.java
index ef49804..0a6b4d0 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/instance/InstanceManager.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/instance/InstanceManager.java
@@ -25,10 +25,12 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 
 import akka.actor.ActorRef;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import scala.Option;
 
 /**
  * Simple manager that keeps track of which TaskManager are available and alive.
@@ -125,7 +127,23 @@ public class InstanceManager {
 		}
 	}
 
-	public InstanceID registerTaskManager(ActorRef taskManager, InstanceConnectionInfo connectionInfo, HardwareDescription resources, int numberOfSlots){
+	/**
+	 * Registers a task manager. Registration of a task manager makes it available to be used
+	 * for the job execution.
+	 *
+	 * @param taskManager ActorRef to the TaskManager which wants to be registered
+	 * @param connectionInfo ConnectionInfo of the TaskManager
+	 * @param resources Hardware description of the TaskManager
+	 * @param numberOfSlots Number of available slots on the TaskManager
+	 * @param leaderSessionID The current leader session ID of the JobManager
+	 * @return
+	 */
+	public InstanceID registerTaskManager(
+			ActorRef taskManager,
+			InstanceConnectionInfo connectionInfo,
+			HardwareDescription resources,
+			int numberOfSlots,
+			Option<UUID> leaderSessionID){
 		synchronized(this.lock){
 			if (this.isShutdown) {
 				throw new IllegalStateException("InstanceManager is shut down.");
@@ -149,9 +167,9 @@ public class InstanceManager {
 				id = new InstanceID();
 			} while (registeredHostsById.containsKey(id));
 
-			InstanceGateway instanceGateway = new AkkaInstanceGateway(taskManager);
+			ActorGateway actorGateway = new AkkaActorGateway(taskManager, leaderSessionID);
 
-			Instance host = new Instance(instanceGateway, connectionInfo, id, resources, numberOfSlots);
+			Instance host = new Instance(actorGateway, connectionInfo, id, resources, numberOfSlots);
 
 			registeredHostsById.put(id, host);
 			registeredHostsByConnection.put(taskManager, host);

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/NetworkEnvironment.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/NetworkEnvironment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/NetworkEnvironment.java
index 0ffc889..539dbc0 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/NetworkEnvironment.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/NetworkEnvironment.java
@@ -18,12 +18,10 @@
 
 package org.apache.flink.runtime.io.network;
 
-import akka.actor.ActorRef;
 import akka.dispatch.OnFailure;
-import akka.pattern.Patterns;
-import akka.util.Timeout;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.io.disk.iomanager.IOManager.IOMode;
 import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
 import org.apache.flink.runtime.io.network.buffer.BufferPool;
@@ -164,16 +162,17 @@ public class NetworkEnvironment {
 	 * This associates the network environment with a TaskManager and JobManager.
 	 * This will actually start the network components.
 	 *
-	 * @param jobManagerRef The JobManager actor reference.
-	 * @param taskManagerRef The TaskManager actor reference.
+	 * @param jobManagerGateway Gateway to the JobManager.
+	 * @param taskManagerGateway Gateway to the TaskManager.
 	 *
 	 * @throws IOException Thrown if the network subsystem (Netty) cannot be properly started.
 	 */
-	public void associateWithTaskManagerAndJobManager(ActorRef jobManagerRef, ActorRef taskManagerRef)
-			throws IOException
+	public void associateWithTaskManagerAndJobManager(
+			ActorGateway jobManagerGateway,
+			ActorGateway taskManagerGateway) throws IOException
 	{
-		checkNotNull(jobManagerRef);
-		checkNotNull(taskManagerRef);
+		checkNotNull(jobManagerGateway);
+		checkNotNull(taskManagerGateway);
 
 		synchronized (lock) {
 			if (isShutdown) {
@@ -192,12 +191,12 @@ public class NetworkEnvironment {
 				this.taskEventDispatcher = new TaskEventDispatcher();
 				this.partitionConsumableNotifier = new JobManagerResultPartitionConsumableNotifier(
 					executionContext,
-					jobManagerRef,
-					taskManagerRef,
-					new Timeout(jobManagerTimeout));
+					jobManagerGateway,
+					taskManagerGateway,
+					jobManagerTimeout);
 
 				this.partitionStateChecker = new JobManagerPartitionStateChecker(
-						jobManagerRef, taskManagerRef);
+						jobManagerGateway, taskManagerGateway);
 
 				// -----  Network connections  -----
 				final Option<NettyConfig> nettyConfig = configuration.nettyConfig();
@@ -432,17 +431,17 @@ public class NetworkEnvironment {
 		 */
 		private final ExecutionContext executionContext;
 
-		private final ActorRef jobManager;
+		private final ActorGateway jobManager;
 
-		private final ActorRef taskManager;
+		private final ActorGateway taskManager;
 
-		private final Timeout jobManagerMessageTimeout;
+		private final FiniteDuration jobManagerMessageTimeout;
 
 		public JobManagerResultPartitionConsumableNotifier(
 			ExecutionContext executionContext,
-			ActorRef jobManager,
-			ActorRef taskManager,
-			Timeout jobManagerMessageTimeout) {
+			ActorGateway jobManager,
+			ActorGateway taskManager,
+			FiniteDuration jobManagerMessageTimeout) {
 
 			this.executionContext = executionContext;
 			this.jobManager = jobManager;
@@ -455,7 +454,7 @@ public class NetworkEnvironment {
 
 			final ScheduleOrUpdateConsumers msg = new ScheduleOrUpdateConsumers(jobId, partitionId);
 
-			Future<Object> futureResponse = Patterns.ask(jobManager, msg, jobManagerMessageTimeout);
+			Future<Object> futureResponse = jobManager.ask(msg, jobManagerMessageTimeout);
 
 			futureResponse.onFailure(new OnFailure() {
 				@Override
@@ -468,7 +467,7 @@ public class NetworkEnvironment {
 							new RuntimeException("Could not notify JobManager to schedule or update consumers",
 									failure));
 
-					taskManager.tell(failMsg, ActorRef.noSender());
+					taskManager.tell(failMsg);
 				}
 			}, executionContext);
 		}
@@ -476,11 +475,11 @@ public class NetworkEnvironment {
 
 	private static class JobManagerPartitionStateChecker implements PartitionStateChecker {
 
-		private final ActorRef jobManager;
+		private final ActorGateway jobManager;
 
-		private final ActorRef taskManager;
+		private final ActorGateway taskManager;
 
-		public JobManagerPartitionStateChecker(ActorRef jobManager, ActorRef taskManager) {
+		public JobManagerPartitionStateChecker(ActorGateway jobManager, ActorGateway taskManager) {
 			this.jobManager = jobManager;
 			this.taskManager = taskManager;
 		}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/web/SetupInfoServlet.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/web/SetupInfoServlet.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/web/SetupInfoServlet.java
index 567d15a..c3df253 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/web/SetupInfoServlet.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/web/SetupInfoServlet.java
@@ -146,7 +146,7 @@ public class SetupInfoServlet extends HttpServlet {
 				long time = new Date().getTime() - instance.getLastHeartBeat();
 
 				try {
-					objInner.put("path", instance.getInstanceGateway().path());
+					objInner.put("path", instance.getActorGateway().path());
 					objInner.put("dataPort", instance.getInstanceConnectionInfo().dataPort());
 					objInner.put("timeSinceLastHeartbeat", time / 1000);
 					objInner.put("slotsNumber", instance.getTotalNumberOfSlots());

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/messages/RequiresLeaderSessionID.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/messages/RequiresLeaderSessionID.java b/flink-runtime/src/main/java/org/apache/flink/runtime/messages/RequiresLeaderSessionID.java
new file mode 100644
index 0000000..7e6a473
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/messages/RequiresLeaderSessionID.java
@@ -0,0 +1,25 @@
+/*
+ * 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.messages;
+
+/**
+ * Marks messages to be sent wrapped in a
+ * {@link org.apache.flink.runtime.messages.JobManagerMessages.LeaderSessionMessage}
+ */
+public interface RequiresLeaderSessionID {}


[4/6] flink git commit: [FLINK-2332] [runtime] Adds leader session IDs and registration session IDs to JobManager and TaskManager messages.

Posted by tr...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/taskmanager/TaskManager.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/taskmanager/TaskManager.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/taskmanager/TaskManager.scala
index d78a594..f974946 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/taskmanager/TaskManager.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/taskmanager/TaskManager.scala
@@ -20,13 +20,14 @@ package org.apache.flink.runtime.taskmanager
 
 import java.io.{File, IOException}
 import java.net.{InetAddress, InetSocketAddress}
+import java.util.UUID
 import java.util.concurrent.TimeUnit
 import java.lang.reflect.Method
 import java.lang.management.{OperatingSystemMXBean, ManagementFactory}
 
-import akka.actor._
-import akka.pattern.ask
-import akka.util.Timeout
+import _root_.akka.actor._
+import _root_.akka.pattern.ask
+import _root_.akka.util.Timeout
 
 import com.codahale.metrics.{Gauge, MetricFilter, MetricRegistry}
 import com.codahale.metrics.json.MetricsModule
@@ -36,9 +37,10 @@ import com.fasterxml.jackson.databind.ObjectMapper
 import grizzled.slf4j.Logger
 
 import org.apache.flink.configuration.{Configuration, ConfigConstants, GlobalConfiguration, IllegalConfigurationException}
+
+import org.apache.flink.runtime.accumulators.AccumulatorSnapshot
 import org.apache.flink.runtime.messages.checkpoint.{NotifyCheckpointComplete, TriggerCheckpoint, AbstractCheckpointMessage}
-import org.apache.flink.runtime.accumulators.{AccumulatorSnapshot, AccumulatorRegistry}
-import org.apache.flink.runtime.{StreamingMode, ActorSynchronousLogging, ActorLogMessages}
+import org.apache.flink.runtime.{FlinkActor, LeaderSessionMessages, LogMessages, StreamingMode}
 import org.apache.flink.runtime.akka.AkkaUtils
 import org.apache.flink.runtime.blob.{BlobService, BlobCache}
 import org.apache.flink.runtime.broadcast.BroadcastVariableManager
@@ -46,7 +48,8 @@ import org.apache.flink.runtime.deployment.{InputChannelDeploymentDescriptor, Ta
 import org.apache.flink.runtime.execution.librarycache.{BlobLibraryCacheManager, FallbackLibraryCacheManager, LibraryCacheManager}
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID
 import org.apache.flink.runtime.filecache.FileCache
-import org.apache.flink.runtime.instance.{HardwareDescription, InstanceConnectionInfo, InstanceID}
+import org.apache.flink.runtime.instance.{AkkaActorGateway, HardwareDescription,
+InstanceConnectionInfo, InstanceID}
 import org.apache.flink.runtime.io.disk.iomanager.IOManager.IOMode
 import org.apache.flink.runtime.io.disk.iomanager.{IOManager, IOManagerAsync}
 import org.apache.flink.runtime.io.network.NetworkEnvironment
@@ -114,15 +117,20 @@ import scala.language.postfixOps
  *    - Exceptions releasing intermediate result resources. Critical resource leak,
  *      requires a clean JVM.
  */
-class TaskManager(protected val config: TaskManagerConfiguration,
-                  protected val connectionInfo: InstanceConnectionInfo,
-                  protected val jobManagerAkkaURL: String,
-                  protected val memoryManager: MemoryManager,
-                  protected val ioManager: IOManager,
-                  protected val network: NetworkEnvironment,
-                  protected val numberOfSlots: Int)
-
-extends Actor with ActorLogMessages with ActorSynchronousLogging {
+class TaskManager(
+    protected val config: TaskManagerConfiguration,
+    protected val connectionInfo: InstanceConnectionInfo,
+    protected val jobManagerAkkaURL: String,
+    protected val memoryManager: MemoryManager,
+    protected val ioManager: IOManager,
+    protected val network: NetworkEnvironment,
+    protected val numberOfSlots: Int)
+  extends FlinkActor
+  with LeaderSessionMessages // Mixin order is important: second we want to filter leader messages
+  with LogMessages // Mixin order is important: first we want to support message logging
+{
+
+  override val log = Logger(getClass)
 
   /** The timeout for all actor ask futures */
   protected val askTimeout = new Timeout(config.timeout)
@@ -161,6 +169,10 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
 
   private var heartbeatScheduler: Option[Cancellable] = None
 
+  protected var leaderSessionID: Option[UUID] = None
+
+  private var currentRegistrationSessionID: UUID = UUID.randomUUID()
+
   // --------------------------------------------------------------------------
   //  Actor messages and life cycle
   // --------------------------------------------------------------------------
@@ -183,8 +195,15 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
     // kick off the registration
     val deadline: Option[Deadline] = config.maxRegistrationDuration.map(_.fromNow)
 
-    self.tell(TriggerTaskManagerRegistration(jobManagerAkkaURL,
-                   TaskManager.INITIAL_REGISTRATION_TIMEOUT, deadline,1), ActorRef.noSender)
+    self ! decorateMessage(
+      TriggerTaskManagerRegistration(
+        currentRegistrationSessionID,
+        jobManagerAkkaURL,
+        TaskManager.INITIAL_REGISTRATION_TIMEOUT,
+        deadline,
+        1)
+    )
+
   }
 
   /**
@@ -236,8 +255,7 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
    * Central handling of actor messages. This method delegates to the more specialized
    * methods for handling certain classes of messages.
    */
-  override def receiveWithLogMessages: Receive = {
-
+  override def handleMessage: Receive = {
     // task messages are most common and critical, we handle them first
     case message: TaskMessage => handleTaskMessage(message)
 
@@ -259,7 +277,7 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
     // its registration at the JobManager
     case NotifyWhenRegisteredAtJobManager =>
       if (isConnected) {
-        sender ! RegisteredAtJobManager
+        sender ! decorateMessage(RegisteredAtJobManager)
       } else {
         waitForRegistration += sender
       }
@@ -339,12 +357,12 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
         // was into a terminal state, or in case the JobManager cannot be informed of the
         // state transition
 
-        case updateMsg@UpdateTaskExecutionState(taskExecutionState: TaskExecutionState) =>
+      case updateMsg@UpdateTaskExecutionState(taskExecutionState: TaskExecutionState) =>
 
           // we receive these from our tasks and forward them to the JobManager
           currentJobManager foreach {
             jobManager => {
-              val futureResponse = (jobManager ? updateMsg)(askTimeout)
+            val futureResponse = (jobManager ? decorateMessage(updateMsg))(askTimeout)
 
               val executionID = taskExecutionState.getID
 
@@ -353,13 +371,20 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
                 //            but only send messages to the TaskManager to do those changes
                 case Success(result) =>
                   if (!result) {
-                    self ! FailTask(executionID,
+                  self ! decorateMessage(
+                    FailTask(
+                      executionID,
                       new Exception("Task has been cancelled on the JobManager."))
+                  )
                   }
 
                 case Failure(t) =>
-                  self ! FailTask(executionID, new Exception(
-                    "Failed to send ExecutionStateChange notification to JobManager"))
+                self ! decorateMessage(
+                  FailTask(
+                    executionID,
+                    new Exception(
+                      "Failed to send ExecutionStateChange notification to JobManager"))
+                )
               }(context.dispatcher)
             }
           }
@@ -387,11 +412,15 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
           val task = runningTasks.get(executionID)
           if (task != null) {
             task.cancelExecution()
-            sender ! new TaskOperationResult(executionID, true)
+          sender ! decorateMessage(new TaskOperationResult(executionID, true))
           } else {
             log.debug(s"Cannot find task to cancel for execution ${executionID})")
-            sender ! new TaskOperationResult(executionID, false,
+          sender ! decorateMessage(
+            new TaskOperationResult(
+              executionID,
+              false,
             "No task with that execution ID was found.")
+          )
           }
 
         case PartitionState(taskExecutionId, taskResultId, partitionId, state) =>
@@ -413,7 +442,6 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
   private def handleCheckpointingMessage(actorMessage: AbstractCheckpointMessage): Unit = {
 
     actorMessage match {
-
       case message: TriggerCheckpoint =>
         val taskExecutionId = message.getTaskExecutionId
         val checkpointId = message.getCheckpointId
@@ -457,118 +485,148 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
    * @param message The registration message.
    */
   private def handleRegistrationMessage(message: RegistrationMessage): Unit = {
+    if(message.registrationSessionID.equals(currentRegistrationSessionID)) {
+      message match {
+        case TriggerTaskManagerRegistration(
+          registrationSessionID,
+          jobManagerURL,
+          timeout,
+          deadline,
+          attempt) =>
+
+          if (isConnected) {
+            // this may be the case, if we queue another attempt and
+            // in the meantime, the registration is acknowledged
+            log.debug(
+              "TaskManager was triggered to register at JobManager, but is already registered")
+          }
+          else if (deadline.exists(_.isOverdue())) {
+            // we failed to register in time. that means we should quit
+            log.error("Failed to register at the JobManager withing the defined maximum " +
+              "connect time. Shutting down ...")
 
-    message match {
-
-      case TriggerTaskManagerRegistration(jobManagerURL, timeout, deadline, attempt) =>
-        if (isConnected) {
-          // this may be the case, if we queue another attempt and
-          // in the meantime, the registration is acknowledged
-          log.debug(
-            "TaskManager was triggered to register at JobManager, but is already registered")
-        }
-        else if (deadline.exists(_.isOverdue())) {
-          // we failed to register in time. that means we should quit
-          log.error("Failed to register at the JobManager withing the defined maximum " +
-            "connect time. Shutting down ...")
+            // terminate ourselves (hasta la vista)
+            self ! decorateMessage(PoisonPill)
+          }
+          else {
+            log.info(s"Trying to register at JobManager ${jobManagerURL} " +
+              s"(attempt ${attempt}, timeout: ${timeout})")
+
+            val jobManager = context.actorSelection(jobManagerAkkaURL)
+            jobManager ! decorateMessage(
+              RegisterTaskManager(
+                registrationSessionID,
+                self,
+                connectionInfo,
+                resources,
+                numberOfSlots)
+            )
+
+            // the next timeout computes via exponential backoff with cap
+            val nextTimeout = (timeout * 2).min(TaskManager.MAX_REGISTRATION_TIMEOUT)
+
+            // schedule (with our timeout s delay) a check triggers a new registration
+            // attempt, if we are not registered by then
+            context.system.scheduler.scheduleOnce(timeout) {
+              if (!isConnected) {
+                self ! decorateMessage(
+                  TriggerTaskManagerRegistration(
+                    registrationSessionID,
+                    jobManagerURL,
+                    nextTimeout,
+                    deadline,
+                    attempt + 1)
+                )
+              }
+            }(context.dispatcher)
+          }
 
-          // terminate ourselves (hasta la vista)
-          self ! PoisonPill
-        }
-        else {
-          log.info(s"Trying to register at JobManager ${jobManagerURL} " +
-            s"(attempt ${attempt}, timeout: ${timeout})")
-
-          val jobManager = context.actorSelection(jobManagerAkkaURL)
-          jobManager ! RegisterTaskManager(self, connectionInfo, resources, numberOfSlots)
-
-          // the next timeout computes via exponential backoff with cap
-          val nextTimeout = (timeout * 2).min(TaskManager.MAX_REGISTRATION_TIMEOUT)
-
-          // schedule (with our timeout s delay) a check triggers a new registration
-          // attempt, if we are not registered by then
-          context.system.scheduler.scheduleOnce(timeout) {
-            if (!isConnected) {
-              self.tell(TriggerTaskManagerRegistration(jobManagerURL,
-                                 nextTimeout, deadline, attempt + 1), ActorRef.noSender)
+        // successful registration. associate with the JobManager
+        // we disambiguate duplicate or erroneous messages, to simplify debugging
+        case AcknowledgeRegistration(_, leaderSessionID, jobManager, id, blobPort) =>
+          if (isConnected) {
+            if (jobManager == currentJobManager.orNull) {
+              log.debug("Ignoring duplicate registration acknowledgement.")
+            } else {
+              log.warn(s"Ignoring 'AcknowledgeRegistration' message from ${jobManager.path} , " +
+                s"because the TaskManager is already registered at ${currentJobManager.orNull}")
             }
-          }(context.dispatcher)
-        }
-
-      // successful registration. associate with the JobManager
-      // we disambiguate duplicate or erroneous messages, to simplify debugging
-      case AcknowledgeRegistration(jobManager, id, blobPort) =>
-        if (isConnected) {
-          if (jobManager == currentJobManager.orNull) {
-            log.debug("Ignoring duplicate registration acknowledgement.")
-          } else {
-            log.warn(s"Ignoring 'AcknowledgeRegistration' message from ${jobManager.path} , " +
-              s"because the TaskManager is already registered at ${currentJobManager.orNull}")
           }
-        }
-        else {
-          // not yet connected, so let's associate with that JobManager
-          try {
-            associateWithJobManager(jobManager, id, blobPort)
-          } catch {
-            case t: Throwable =>
-              killTaskManagerFatal(
-                "Unable to start TaskManager components after registering at JobManager", t)
+          else {
+            // not yet connected, so let's associate with that JobManager
+            try {
+              associateWithJobManager(jobManager, id, blobPort, leaderSessionID)
+            } catch {
+              case t: Throwable =>
+                killTaskManagerFatal(
+                  "Unable to start TaskManager components after registering at JobManager", t)
+            }
           }
-        }
 
-      // we are already registered at that specific JobManager - duplicate answer, rare cases
-      case AlreadyRegistered(jobManager, id, blobPort) =>
-        if (isConnected) {
-          if (jobManager == currentJobManager.orNull) {
-            log.debug("Ignoring duplicate registration acknowledgement.")
-          } else {
-            log.warn(s"Received 'AlreadyRegistered' message from JobManager ${jobManager.path}, " +
-              s"even through TaskManager is currently registered at ${currentJobManager.orNull}")
+        // we are already registered at that specific JobManager - duplicate answer, rare cases
+        case AlreadyRegistered(_, leaderSesssionID, jobManager, id, blobPort) =>
+          if (isConnected) {
+            if (jobManager == currentJobManager.orNull) {
+              log.debug("Ignoring duplicate registration acknowledgement.")
+            } else {
+              log.warn(s"Received 'AlreadyRegistered' message from " +
+                s"JobManager ${jobManager.path}, even through TaskManager is currently " +
+                s"registered at ${currentJobManager.orNull}")
+            }
           }
-        }
-        else {
-          // not connected, yet, to let's associate
-          log.info("Received 'AlreadyRegistered' message before 'AcknowledgeRegistration'")
-
-          try {
-            associateWithJobManager(jobManager, id, blobPort)
-          } catch {
-            case t: Throwable =>
-              killTaskManagerFatal(
-                "Unable to start TaskManager components after registering at JobManager", t)
+          else {
+            // not connected, yet, to let's associate
+            log.info("Received 'AlreadyRegistered' message before 'AcknowledgeRegistration'")
+
+            try {
+              associateWithJobManager(jobManager, id, blobPort, leaderSesssionID)
+            } catch {
+              case t: Throwable =>
+                killTaskManagerFatal(
+                  "Unable to start TaskManager components after registering at JobManager", t)
+            }
           }
-        }
 
-      case RefuseRegistration(reason) =>
-        if (currentJobManager.isEmpty) {
-          log.error(s"The registration at JobManager ${jobManagerAkkaURL} was refused, " +
-                    s"because: ${reason}. Retrying later...")
+        case RefuseRegistration(registrationSessionID, reason) =>
+          if (currentJobManager.isEmpty) {
+            log.error(s"The registration at JobManager ${jobManagerAkkaURL} was refused, " +
+              s"because: ${reason}. Retrying later...")
 
-          // try the registration again after some time
+            // try the registration again after some time
 
-          val delay: FiniteDuration = TaskManager.DELAY_AFTER_REFUSED_REGISTRATION
-          val deadline: Option[Deadline] = config.maxRegistrationDuration.map {
-            timeout => timeout + delay fromNow
-          }
+            val delay: FiniteDuration = TaskManager.DELAY_AFTER_REFUSED_REGISTRATION
+            val deadline: Option[Deadline] = config.maxRegistrationDuration.map {
+              timeout => timeout + delay fromNow
+            }
 
-          context.system.scheduler.scheduleOnce(delay) {
-            self.tell(TriggerTaskManagerRegistration(jobManagerAkkaURL,
-                  TaskManager.INITIAL_REGISTRATION_TIMEOUT, deadline, 1), ActorRef.noSender)
-          }(context.dispatcher)
-        }
-        else {
-          // ignore RefuseRegistration messages which arrived after AcknowledgeRegistration
-          if (sender() == currentJobManager.orNull) {
-            log.warn(s"Received 'RefuseRegistration' from the JobManager (${sender().path})" +
-                     s" even though this TaskManager is already registered there.")
+            context.system.scheduler.scheduleOnce(delay) {
+              self ! decorateMessage(
+                TriggerTaskManagerRegistration(
+                  registrationSessionID,
+                  jobManagerAkkaURL,
+                  TaskManager.INITIAL_REGISTRATION_TIMEOUT,
+                  deadline,
+                  1)
+              )
+            }(context.dispatcher)
           }
           else {
-            log.warn(s"Ignoring 'RefuseRegistration' from unrelated JobManager (${sender().path})")
+            // ignore RefuseRegistration messages which arrived after AcknowledgeRegistration
+            if (sender() == currentJobManager.orNull) {
+              log.warn(s"Received 'RefuseRegistration' from the JobManager (${sender().path})" +
+                s" even though this TaskManager is already registered there.")
+            }
+            else {
+              log.warn(s"Ignoring 'RefuseRegistration' from unrelated " +
+                s"JobManager (${sender().path})")
+            }
           }
-        }
 
-      case _ => unhandled(message)
+        case _ => unhandled(message)
+      }
+    } else {
+      log.debug(s"Discarded registration message ${message}, because the registration session " +
+        "ID was not correct.")
     }
   }
 
@@ -593,20 +651,28 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
    * @param id The instanceID under which the TaskManager is registered
    *           at the JobManager.
    * @param blobPort The JobManager's port for the BLOB server.
+   * @param newLeaderSessionID Leader session ID of the JobManager
    */
-  private def associateWithJobManager(jobManager: ActorRef,
-                                      id: InstanceID,
-                                      blobPort: Int): Unit = {
+  private def associateWithJobManager(
+      jobManager: ActorRef,
+      id: InstanceID,
+      blobPort: Int,
+      newLeaderSessionID: UUID)
+    : Unit = {
 
     if (jobManager == null) {
-      throw new NullPointerException("jobManager may not be null")
+      throw new NullPointerException("jobManager must not be null.")
     }
     if (id == null) {
-      throw new NullPointerException("instance ID may not be null")
+      throw new NullPointerException("instance ID must not be null.")
     }
     if (blobPort <= 0 || blobPort > 65535) {
       throw new IllegalArgumentException("blob port is out of range: " + blobPort)
     }
+    if(newLeaderSessionID == null) {
+      throw new NullPointerException("Leader session ID must not be null.")
+    }
+
 
     // sanity check that we are not currently registered with a different JobManager
     if (isConnected) {
@@ -631,9 +697,18 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
       throw new IllegalStateException("JobManager-specific components are already initialized.")
     }
 
+    currentJobManager = Some(jobManager)
+    instanceID = id
+    leaderSessionID = Some(newLeaderSessionID)
+
     // start the network stack, now that we have the JobManager actor reference
     try {
-      network.associateWithTaskManagerAndJobManager(jobManager, self)
+      network.associateWithTaskManagerAndJobManager(
+        new AkkaActorGateway(jobManager, leaderSessionID),
+        new AkkaActorGateway(self, leaderSessionID)
+      )
+
+
     }
     catch {
       case e: Exception =>
@@ -665,16 +740,18 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
       libraryCacheManager = Some(new FallbackLibraryCacheManager)
     }
 
-    currentJobManager = Some(jobManager)
-    instanceID = id
-
     // watch job manager to detect when it dies
     context.watch(jobManager)
 
     // schedule regular heartbeat message for oneself
-    heartbeatScheduler = Some(context.system.scheduler.schedule(
-      TaskManager.HEARTBEAT_INTERVAL, TaskManager.HEARTBEAT_INTERVAL, self, SendHeartbeat)
-       (context.dispatcher))
+    heartbeatScheduler = Some(
+      context.system.scheduler.schedule(
+        TaskManager.HEARTBEAT_INTERVAL,
+        TaskManager.HEARTBEAT_INTERVAL,
+        self,
+        decorateMessage(SendHeartbeat)
+      )(context.dispatcher)
+    )
 
     // notify all the actors that listen for a successful registration
     for (listener <- waitForRegistration) {
@@ -710,7 +787,7 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
 
     // de-register from the JobManager (faster detection of disconnect)
     currentJobManager foreach {
-      _ ! Disconnect(s"TaskManager ${self.path} is shutting down.")
+      _ ! decorateMessage(Disconnect(s"TaskManager ${self.path} is shutting down."))
     }
 
     currentJobManager = None
@@ -748,8 +825,14 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
 
           // begin attempts to reconnect
           val deadline: Option[Deadline] = config.maxRegistrationDuration.map(_.fromNow)
-          self ! TriggerTaskManagerRegistration(jobManagerAkkaURL,
-                               TaskManager.INITIAL_REGISTRATION_TIMEOUT, deadline, 1)
+          self ! decorateMessage(
+            TriggerTaskManagerRegistration(
+              currentRegistrationSessionID,
+              jobManagerAkkaURL,
+              TaskManager.INITIAL_REGISTRATION_TIMEOUT,
+              deadline,
+              1)
+          )
         }
         catch {
           // this is pretty bad, it leaves the TaskManager in a state where it cannot
@@ -795,8 +878,21 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
 
       // create the task. this does not grab any TaskManager resources or download
       // and libraries - the operation does not block
-      val task = new Task(tdd, memoryManager, ioManager, network, bcVarManager,
-                          self, jobManagerActor, config.timeout, libCache, fileCache)
+
+      val jobManagerGateway = new AkkaActorGateway(jobManagerActor, leaderSessionID)
+      val selfGateway = new AkkaActorGateway(self, leaderSessionID)
+
+      val task = new Task(
+        tdd,
+        memoryManager,
+        ioManager,
+        network,
+        bcVarManager,
+        selfGateway,
+        jobManagerGateway,
+        config.timeout,
+        libCache,
+        fileCache)
 
       log.info(s"Received task ${task.getTaskNameWithSubtasks}")
 
@@ -812,12 +908,12 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
       // all good, we kick off the task, which performs its own initialization
       task.startTaskThread()
       
-      sender ! Acknowledge
+      sender ! decorateMessage(Acknowledge)
     }
     catch {
       case t: Throwable => 
         log.error("SubmitTask failed", t)
-        sender ! Failure(t)
+        sender ! decorateMessage(Failure(t))
     }
   }
 
@@ -828,8 +924,9 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
    * @param partitionInfos The descriptor of the intermediate result partitions.
    */
   private def updateTaskInputPartitions(
-         executionId: ExecutionAttemptID,
-         partitionInfos: Seq[(IntermediateDataSetID, InputChannelDeploymentDescriptor)]) : Unit = {
+       executionId: ExecutionAttemptID,
+       partitionInfos: Seq[(IntermediateDataSetID, InputChannelDeploymentDescriptor)])
+    : Unit = {
 
     Option(runningTasks.get(executionId)) match {
       case Some(task) =>
@@ -867,15 +964,15 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
         }
 
         if (errors.isEmpty) {
-          sender ! Acknowledge
+          sender ! decorateMessage(Acknowledge)
         } else {
-          sender ! Failure(new Exception(errors.mkString("\n")))
+          sender ! decorateMessage(Failure(new Exception(errors.mkString("\n"))))
         }
 
       case None =>
         log.debug(s"Discard update for input partitions of task $executionId : " +
           s"task is no longer running.")
-        sender ! Acknowledge
+        sender ! decorateMessage(Acknowledge)
     }
   }
 
@@ -919,9 +1016,16 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
         registry.getSnapshot
       }
 
-      self ! UpdateTaskExecutionState(new TaskExecutionState(
-        task.getJobID, task.getExecutionId, task.getExecutionState, task.getFailureCause,
-        accumulators))
+        self ! decorateMessage(
+          UpdateTaskExecutionState(
+            new TaskExecutionState(
+              task.getJobID,
+              task.getExecutionId,
+              task.getExecutionState,
+              task.getFailureCause,
+              accumulators)
+          )
+        )
     }
     else {
       log.error(s"Cannot find task with ID $executionID to unregister.")
@@ -952,7 +1056,7 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
       }
 
        currentJobManager foreach {
-        jm => jm ! Heartbeat(instanceID, metricsReport, accumulatorEvents)
+        jm => jm ! decorateMessage(Heartbeat(instanceID, metricsReport, accumulatorEvents))
       }
     }
     catch {
@@ -972,14 +1076,12 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
 
     try {
       val traces = Thread.getAllStackTraces.asScala
-      val stackTraceStr = traces.map(
-        (trace: (Thread, Array[StackTraceElement])) => {
-          val (thread, elements) = trace
+      val stackTraceStr = traces.map {
+        case (thread: Thread, elements: Array[StackTraceElement]) =>
           "Thread: " + thread.getName + '\n' + elements.mkString("\n")
-          })
-        .mkString("\n\n")
+        }.mkString("\n\n")
 
-      recipient ! StackTrace(instanceID, stackTraceStr)
+      recipient ! decorateMessage(StackTrace(instanceID, stackTraceStr))
     }
     catch {
       case e: Exception => log.error("Failed to send stack trace to " + recipient.path, e)
@@ -999,7 +1101,7 @@ extends Actor with ActorLogMessages with ActorSynchronousLogging {
       "\n" +
       "A fatal error occurred, forcing the TaskManager to shut down: " + message, cause)
 
-    self ! Kill
+    self ! decorateMessage(Kill)
   }
 }
 
@@ -1165,24 +1267,34 @@ object TaskManager {
    *                         Allows to use TaskManager subclasses for example for YARN.
    */
   @throws(classOf[Exception])
-  def selectNetworkInterfaceAndRunTaskManager(configuration: Configuration,
-                                              streamingMode: StreamingMode,
-                                              taskManagerClass: Class[_ <: TaskManager]) : Unit = {
+  def selectNetworkInterfaceAndRunTaskManager(
+      configuration: Configuration,
+      streamingMode: StreamingMode,
+      taskManagerClass: Class[_ <: TaskManager])
+    : Unit = {
 
     val (jobManagerHostname, jobManagerPort) = getAndCheckJobManagerAddress(configuration)
 
-    val (taskManagerHostname, actorSystemPort) =
-       selectNetworkInterfaceAndPort(configuration, jobManagerHostname, jobManagerPort)
+    val (taskManagerHostname, actorSystemPort) = selectNetworkInterfaceAndPort(
+      configuration,
+      jobManagerHostname,
+      jobManagerPort)
 
-    runTaskManager(taskManagerHostname, actorSystemPort, configuration,
-                   streamingMode, taskManagerClass)
+    runTaskManager(
+      taskManagerHostname,
+      actorSystemPort,
+      configuration,
+      streamingMode,
+      taskManagerClass)
   }
 
   @throws(classOf[IOException])
   @throws(classOf[IllegalConfigurationException])
-  def selectNetworkInterfaceAndPort(configuration: Configuration,
-                                    jobManagerHostname: String,
-                                    jobManagerPort: Int) : (String, Int) = {
+  def selectNetworkInterfaceAndPort(
+      configuration: Configuration,
+      jobManagerHostname: String,
+      jobManagerPort: Int)
+    : (String, Int) = {
 
     var taskManagerHostname = configuration.getString(
       ConfigConstants.TASK_MANAGER_HOSTNAME_KEY, null)
@@ -1243,13 +1355,19 @@ object TaskManager {
    * @param configuration The configuration for the TaskManager.
    */
   @throws(classOf[Exception])
-  def runTaskManager(taskManagerHostname: String,
-                     actorSystemPort: Int, 
-                     configuration: Configuration,
-                     streamingMode: StreamingMode) : Unit = {
+  def runTaskManager(
+      taskManagerHostname: String,
+      actorSystemPort: Int,
+      configuration: Configuration,
+      streamingMode: StreamingMode)
+    : Unit = {
 
-    runTaskManager(taskManagerHostname, actorSystemPort, configuration,
-                   streamingMode, classOf[TaskManager])
+    runTaskManager(
+      taskManagerHostname,
+      actorSystemPort,
+      configuration,
+      streamingMode,
+      classOf[TaskManager])
   }
 
   /**
@@ -1269,11 +1387,13 @@ object TaskManager {
    *                         subclasses for example for YARN.
    */
   @throws(classOf[Exception])
-  def runTaskManager(taskManagerHostname: String,
-                     actorSystemPort: Int,
-                     configuration: Configuration,
-                     streamingMode: StreamingMode,
-                     taskManagerClass: Class[_ <: TaskManager]) : Unit = {
+  def runTaskManager(
+      taskManagerHostname: String,
+      actorSystemPort: Int,
+      configuration: Configuration,
+      streamingMode: StreamingMode,
+      taskManagerClass: Class[_ <: TaskManager])
+    : Unit = {
 
     LOG.info(s"Starting TaskManager in streaming mode $streamingMode")
 
@@ -1282,8 +1402,10 @@ object TaskManager {
     LOG.info(s"Starting TaskManager actor system at $taskManagerHostname:$actorSystemPort")
 
     val taskManagerSystem = try {
-      val akkaConfig = AkkaUtils.getAkkaConfig(configuration,
-                                               Some((taskManagerHostname, actorSystemPort)))
+      val akkaConfig = AkkaUtils.getAkkaConfig(
+        configuration,
+        Some((taskManagerHostname, actorSystemPort))
+      )
       if (LOG.isDebugEnabled) {
         LOG.debug("Using akka configuration\n " + akkaConfig)
       }
@@ -1307,13 +1429,14 @@ object TaskManager {
     // and the TaskManager actor
     try {
       LOG.info("Starting TaskManager actor")
-      val taskManager = startTaskManagerComponentsAndActor(configuration,
-                                                           taskManagerSystem,
-                                                           taskManagerHostname,
-                                                           Some(TASK_MANAGER_NAME),
-                                                           None, false,
-                                                           streamingMode,
-                                                           taskManagerClass)
+      val taskManager = startTaskManagerComponentsAndActor(
+        configuration,
+        taskManagerSystem,
+        taskManagerHostname,
+        Some(TASK_MANAGER_NAME),
+        None, false,
+        streamingMode,
+        taskManagerClass)
 
       // start a process reaper that watches the JobManager. If the TaskManager actor dies,
       // the process reaper will kill the JVM process (to ensure easy failure detection)
@@ -1433,14 +1556,15 @@ object TaskManager {
       configuredMemory << 20 // megabytes to bytes
     }
     else {
-      val fraction = configuration.getFloat(ConfigConstants.TASK_MANAGER_MEMORY_FRACTION_KEY,
-                                            ConfigConstants.DEFAULT_MEMORY_MANAGER_MEMORY_FRACTION)
+      val fraction = configuration.getFloat(
+        ConfigConstants.TASK_MANAGER_MEMORY_FRACTION_KEY,
+        ConfigConstants.DEFAULT_MEMORY_MANAGER_MEMORY_FRACTION)
       checkConfigParameter(fraction > 0.0f && fraction < 1.0f, fraction,
                            ConfigConstants.TASK_MANAGER_MEMORY_FRACTION_KEY,
                            "MemoryManager fraction of the free memory must be between 0.0 and 1.0")
 
       val relativeMemSize = (EnvironmentInformation.getSizeOfFreeHeapMemoryWithDefrag() *
-                                                                                   fraction).toLong
+        fraction).toLong
 
       LOG.info(s"Using $fraction of the currently free heap space for Flink managed " +
         s"memory (${relativeMemSize >> 20} MB).")
@@ -1452,10 +1576,11 @@ object TaskManager {
 
     // now start the memory manager
     val memoryManager = try {
-      new DefaultMemoryManager(memorySize,
-                               taskManagerConfig.numberOfSlots,
-                               netConfig.networkBufferSize,
-                               preAllocateMemory)
+      new DefaultMemoryManager(
+        memorySize,
+        taskManagerConfig.numberOfSlots,
+        netConfig.networkBufferSize,
+        preAllocateMemory)
     }
     catch {
       case e: OutOfMemoryError => throw new Exception(
@@ -1472,7 +1597,8 @@ object TaskManager {
     }
 
     // create the actor properties (which define the actor constructor parameters)
-    val tmProps = Props(taskManagerClass,
+    val tmProps = Props(
+      taskManagerClass,
       taskManagerConfig,
       connectionInfo,
       jobManagerAkkaUrl,
@@ -1502,9 +1628,11 @@ object TaskManager {
    * @return The ActorRef to the TaskManager
    */
   @throws(classOf[IOException])
-  def getTaskManagerRemoteReference(taskManagerUrl: String,
-                                    system: ActorSystem,
-                                    timeout: FiniteDuration): ActorRef = {
+  def getTaskManagerRemoteReference(
+      taskManagerUrl: String,
+      system: ActorSystem,
+      timeout: FiniteDuration)
+    : ActorRef = {
     try {
       val future = AkkaUtils.getReference(taskManagerUrl, system, timeout)
       Await.result(future, timeout)
@@ -1536,10 +1664,11 @@ object TaskManager {
    *                  InstanceConnectionInfo, JobManager actor Akka URL).
    */
   @throws(classOf[IllegalArgumentException])
-  def parseTaskManagerConfiguration(configuration: Configuration,
-                                    taskManagerHostname: String,
-                                    localTaskManagerCommunication: Boolean):
-    (TaskManagerConfiguration,
+  def parseTaskManagerConfiguration(
+      configuration: Configuration,
+      taskManagerHostname: String,
+      localTaskManagerCommunication: Boolean)
+    : (TaskManagerConfiguration,
      NetworkEnvironmentConfiguration,
      InstanceConnectionInfo) = {
 
@@ -1579,10 +1708,10 @@ object TaskManager {
       ConfigConstants.TASK_MANAGER_NETWORK_NUM_BUFFERS_KEY)
     
     val pageSizeNew: Int = configuration.getInteger(
-                                        ConfigConstants.TASK_MANAGER_MEMORY_SEGMENT_SIZE_KEY, -1)
+      ConfigConstants.TASK_MANAGER_MEMORY_SEGMENT_SIZE_KEY, -1)
     
     val pageSizeOld: Int = configuration.getInteger(
-                                        ConfigConstants.TASK_MANAGER_NETWORK_BUFFER_SIZE_KEY, -1)
+      ConfigConstants.TASK_MANAGER_NETWORK_BUFFER_SIZE_KEY, -1)
 
     val pageSize: Int =
       if (pageSizeNew != -1) {
@@ -1617,7 +1746,7 @@ object TaskManager {
     val tmpDirs = configuration.getString(
       ConfigConstants.TASK_MANAGER_TMP_DIR_KEY,
       ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH)
-      .split(",|" + File.pathSeparator)
+    .split(",|" + File.pathSeparator)
 
     val nettyConfig = if (localTaskManagerCommunication) {
       None
@@ -1633,7 +1762,10 @@ object TaskManager {
     val ioMode : IOMode = if (syncOrAsync == "async") IOMode.ASYNC else IOMode.SYNC
 
     val networkConfig = NetworkEnvironmentConfiguration(
-      numNetworkBuffers, pageSize, ioMode, nettyConfig)
+      numNetworkBuffers,
+      pageSize,
+      ioMode,
+      nettyConfig)
 
     // ----> timeouts, library caching, profiling
 
@@ -1667,8 +1799,12 @@ object TaskManager {
         e)
     }
 
-    val taskManagerConfig = TaskManagerConfiguration(tmpDirs, cleanupInterval, timeout,
-      finiteRegistratioDuration, slots,
+    val taskManagerConfig = TaskManagerConfiguration(
+      tmpDirs,
+      cleanupInterval,
+      timeout,
+      finiteRegistratioDuration,
+      slots,
       configuration)
 
     (taskManagerConfig, networkConfig, connectionInfo)
@@ -1720,10 +1856,12 @@ object TaskManager {
    * @throws IllegalConfigurationException Thrown if the condition is violated.
    */
   @throws(classOf[IllegalConfigurationException])
-  private def checkConfigParameter(condition: Boolean,
-                                   parameter: Any,
-                                   name: String,
-                                   errorMessage: String = ""): Unit = {
+  private def checkConfigParameter(
+      condition: Boolean,
+      parameter: Any,
+      name: String,
+      errorMessage: String = "")
+    : Unit = {
     if (!condition) {
       throw new IllegalConfigurationException(
         s"Invalid configuration value for '${name}' : ${parameter} - ${errorMessage}")

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/akka/FlinkUntypedActorTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/akka/FlinkUntypedActorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/akka/FlinkUntypedActorTest.java
new file mode 100644
index 0000000..324b014
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/akka/FlinkUntypedActorTest.java
@@ -0,0 +1,146 @@
+/*
+ * 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.akka;
+
+import akka.actor.ActorRef;
+import akka.actor.ActorSystem;
+import akka.actor.Kill;
+import akka.actor.Props;
+import akka.testkit.JavaTestKit;
+import akka.testkit.TestActorRef;
+import org.apache.flink.runtime.messages.JobManagerMessages;
+import org.apache.flink.runtime.messages.RequiresLeaderSessionID;
+import org.apache.flink.runtime.testingUtils.TestingUtils;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import scala.Option;
+
+import java.util.UUID;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+public class FlinkUntypedActorTest {
+
+	private static ActorSystem actorSystem;
+
+	@BeforeClass
+	public static void setup() {
+		actorSystem = ActorSystem.create("TestingActorSystem", TestingUtils.testConfig());
+	}
+
+	@AfterClass
+	public static void teardown() {
+		JavaTestKit.shutdownActorSystem(actorSystem);
+	}
+
+	@Test
+	public void testLeaderSessionMessageFilteringOfFlinkUntypedActor() {
+		final Option<UUID> leaderSessionID = Option.apply(UUID.randomUUID());
+		final Option<UUID> oldSessionID = Option.apply(UUID.randomUUID());
+
+		TestActorRef<PlainFlinkUntypedActor> actor = null;
+
+		try {
+			actor = TestActorRef.create(
+					actorSystem, Props.create(PlainFlinkUntypedActor.class, leaderSessionID));
+
+			final PlainFlinkUntypedActor underlyingActor = actor.underlyingActor();
+
+			actor.tell(new JobManagerMessages.LeaderSessionMessage(leaderSessionID, 1), ActorRef.noSender());
+			actor.tell(new JobManagerMessages.LeaderSessionMessage(oldSessionID, 2), ActorRef.noSender());
+			actor.tell(new JobManagerMessages.LeaderSessionMessage(leaderSessionID, 2), ActorRef.noSender());
+			actor.tell(1, ActorRef.noSender());
+
+			assertEquals(3, underlyingActor.getMessageCounter());
+
+		} finally {
+			stopActor(actor);
+		}
+	}
+
+	@Test
+	public void testThrowingExceptionWhenReceivingNonWrappedRequiresLeaderSessionIDMessage() {
+		final Option<UUID> leaderSessionID = Option.apply(UUID.randomUUID());
+
+		TestActorRef<PlainFlinkUntypedActor> actor = null;
+
+		try{
+			final Props props = Props.create(PlainFlinkUntypedActor.class, leaderSessionID);
+			actor = TestActorRef.create(actorSystem, props);
+
+			actor.receive(new JobManagerMessages.LeaderSessionMessage(leaderSessionID, 1));
+
+			try {
+				actor.receive(new PlainRequiresLeaderSessionID());
+
+				fail("Expected an exception to be thrown, because a RequiresLeaderSessionID" +
+						"message was sent without being wrapped in LeaderSessionMessage.");
+			} catch (Exception e) {
+				assertEquals("Received a message PlainRequiresLeaderSessionID " +
+						"without a leader session ID, even though it requires to have one.",
+						e.getMessage());
+			}
+
+		} finally {
+			stopActor(actor);
+		}
+	}
+
+	private static void stopActor(ActorRef actor) {
+		if(actor != null) {
+			actor.tell(Kill.getInstance(), ActorRef.noSender());
+		}
+	}
+
+
+	static class PlainFlinkUntypedActor extends FlinkUntypedActor {
+
+		private Option<UUID> leaderSessionID;
+
+		private int messageCounter;
+
+		public PlainFlinkUntypedActor(Option<UUID> leaderSessionID) {
+			this.leaderSessionID = leaderSessionID;
+			this.messageCounter = 0;
+		}
+
+		@Override
+		protected void handleMessage(Object message) throws Exception {
+			messageCounter++;
+		}
+
+		@Override
+		protected Option<UUID> getLeaderSessionID() {
+			return leaderSessionID;
+		}
+
+		public int getMessageCounter() {
+			return messageCounter;
+		}
+	}
+
+	static class PlainRequiresLeaderSessionID implements RequiresLeaderSessionID {
+		@Override
+		public String toString() {
+			return "PlainRequiresLeaderSessionID";
+		}
+	}
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CoordinatorShutdownTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CoordinatorShutdownTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CoordinatorShutdownTest.java
index 932e366..b124304 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CoordinatorShutdownTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CoordinatorShutdownTest.java
@@ -18,12 +18,10 @@
 
 package org.apache.flink.runtime.checkpoint;
 
-import akka.actor.ActorRef;
-import akka.pattern.Patterns;
-
 import org.apache.flink.configuration.ConfigConstants;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.executiongraph.ExecutionGraph;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.jobgraph.JobVertex;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
@@ -62,18 +60,19 @@ public class CoordinatorShutdownTest {
 			JobGraph testGraph = new JobGraph("test job", vertex);
 			testGraph.setSnapshotSettings(new JobSnapshottingSettings(vertexIdList, vertexIdList, vertexIdList, 5000));
 			
-			ActorRef jobManager = cluster.getJobManager();
+			ActorGateway jobManager = cluster.getJobManagerGateway();
 
 			FiniteDuration timeout = new FiniteDuration(60, TimeUnit.SECONDS);
 			JobManagerMessages.SubmitJob submitMessage = new JobManagerMessages.SubmitJob(testGraph, false);
 			
 			// submit is successful, but then the job dies because no TaskManager / slot is available
-			Future<Object> submitFuture = Patterns.ask(jobManager, submitMessage, timeout.toMillis());
+			Future<Object> submitFuture = jobManager.ask(submitMessage, timeout);
 			Await.result(submitFuture, timeout);
 
 			// get the execution graph and make sure the coordinator is properly shut down
-			Future<Object> jobRequestFuture = Patterns.ask(jobManager,
-					new JobManagerMessages.RequestJob(testGraph.getJobID()), timeout.toMillis());
+			Future<Object> jobRequestFuture = jobManager.ask(
+					new JobManagerMessages.RequestJob(testGraph.getJobID()),
+					timeout);
 			
 			ExecutionGraph graph = ((JobManagerMessages.JobFound) Await.result(jobRequestFuture, timeout)).executionGraph();
 			
@@ -109,18 +108,19 @@ public class CoordinatorShutdownTest {
 			JobGraph testGraph = new JobGraph("test job", vertex);
 			testGraph.setSnapshotSettings(new JobSnapshottingSettings(vertexIdList, vertexIdList, vertexIdList, 5000));
 			
-			ActorRef jobManager = cluster.getJobManager();
+			ActorGateway jobManager = cluster.getJobManagerGateway();
 
 			FiniteDuration timeout = new FiniteDuration(60, TimeUnit.SECONDS);
 			JobManagerMessages.SubmitJob submitMessage = new JobManagerMessages.SubmitJob(testGraph, false);
 
 			// submit is successful, but then the job dies because no TaskManager / slot is available
-			Future<Object> submitFuture = Patterns.ask(jobManager, submitMessage, timeout.toMillis());
+			Future<Object> submitFuture = jobManager.ask(submitMessage, timeout);
 			Await.result(submitFuture, timeout);
 
 			// get the execution graph and make sure the coordinator is properly shut down
-			Future<Object> jobRequestFuture = Patterns.ask(jobManager,
-					new JobManagerMessages.RequestJob(testGraph.getJobID()), timeout.toMillis());
+			Future<Object> jobRequestFuture = jobManager.ask(
+					new JobManagerMessages.RequestJob(testGraph.getJobID()),
+					timeout);
 
 			ExecutionGraph graph = ((JobManagerMessages.JobFound) Await.result(jobRequestFuture, timeout)).executionGraph();
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphDeploymentTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphDeploymentTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphDeploymentTest.java
index cff7146..e3fc852 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphDeploymentTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphDeploymentTest.java
@@ -92,7 +92,7 @@ public class ExecutionGraphDeploymentTest {
 			ExecutionJobVertex ejv = eg.getAllVertices().get(jid2);
 			ExecutionVertex vertex = ejv.getTaskVertices()[3];
 
-			ExecutionGraphTestUtils.SimpleInstanceGateway instanceGateway = new ExecutionGraphTestUtils.SimpleInstanceGateway(TestingUtils.directExecutionContext());
+			ExecutionGraphTestUtils.SimpleActorGateway instanceGateway = new ExecutionGraphTestUtils.SimpleActorGateway(TestingUtils.directExecutionContext());
 
 			final Instance instance = getInstance(instanceGateway);
 
@@ -295,7 +295,7 @@ public class ExecutionGraphDeploymentTest {
 		for (int i = 0; i < dop1 + dop2; i++) {
 			scheduler.newInstanceAvailable(
 					ExecutionGraphTestUtils.getInstance(
-							new ExecutionGraphTestUtils.SimpleInstanceGateway(
+							new ExecutionGraphTestUtils.SimpleActorGateway(
 									TestingUtils.directExecutionContext())));
 		}
 		assertEquals(dop1 + dop2, scheduler.getNumberOfAvailableSlots());

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphTestUtils.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphTestUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphTestUtils.java
index 8a63060..64d4c44 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphTestUtils.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphTestUtils.java
@@ -24,18 +24,17 @@ import static org.mockito.Mockito.spy;
 
 import java.lang.reflect.Field;
 import java.net.InetAddress;
-import java.util.LinkedList;
 
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.JobException;
 import org.apache.flink.runtime.akka.AkkaUtils;
 import org.apache.flink.runtime.deployment.TaskDeploymentDescriptor;
 import org.apache.flink.runtime.execution.ExecutionState;
-import org.apache.flink.runtime.instance.BaseTestingInstanceGateway;
+import org.apache.flink.runtime.instance.BaseTestingActorGateway;
 import org.apache.flink.runtime.instance.HardwareDescription;
 import org.apache.flink.runtime.instance.Instance;
 import org.apache.flink.runtime.instance.InstanceConnectionInfo;
-import org.apache.flink.runtime.instance.InstanceGateway;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.instance.InstanceID;
 import org.apache.flink.runtime.instance.SimpleSlot;
 import org.apache.flink.runtime.jobgraph.JobVertex;
@@ -101,11 +100,11 @@ public class ExecutionGraphTestUtils {
 	//  utility mocking methods
 	// --------------------------------------------------------------------------------------------
 
-	public static Instance getInstance(final InstanceGateway gateway) throws Exception {
+	public static Instance getInstance(final ActorGateway gateway) throws Exception {
 		return getInstance(gateway, 1);
 	}
 
-	public static Instance getInstance(final InstanceGateway gateway, final int numberOfSlots) throws Exception {
+	public static Instance getInstance(final ActorGateway gateway, final int numberOfSlots) throws Exception {
 		HardwareDescription hardwareDescription = new HardwareDescription(4, 2L*1024*1024*1024, 1024*1024*1024, 512*1024*1024);
 		InetAddress address = InetAddress.getByName("127.0.0.1");
 		InstanceConnectionInfo connection = new InstanceConnectionInfo(address, 10001);
@@ -113,10 +112,10 @@ public class ExecutionGraphTestUtils {
 		return new Instance(gateway, connection, new InstanceID(), hardwareDescription, numberOfSlots);
 	}
 
-	public static class SimpleInstanceGateway extends BaseTestingInstanceGateway {
+	public static class SimpleActorGateway extends BaseTestingActorGateway {
 		public TaskDeploymentDescriptor lastTDD;
 
-		public SimpleInstanceGateway(ExecutionContext executionContext){
+		public SimpleActorGateway(ExecutionContext executionContext){
 			super(executionContext);
 		}
 
@@ -140,8 +139,8 @@ public class ExecutionGraphTestUtils {
 		}
 	}
 
-	public static class SimpleFailingInstanceGateway extends BaseTestingInstanceGateway {
-		public SimpleFailingInstanceGateway(ExecutionContext executionContext) {
+	public static class SimpleFailingActorGateway extends BaseTestingActorGateway {
+		public SimpleFailingActorGateway(ExecutionContext executionContext) {
 			super(executionContext);
 		}
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionStateProgressTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionStateProgressTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionStateProgressTest.java
index f47e92c..89b82f3 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionStateProgressTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionStateProgressTest.java
@@ -63,7 +63,7 @@ public class ExecutionStateProgressTest {
 			// mock resources and mock taskmanager
 			for (ExecutionVertex ee : ejv.getTaskVertices()) {
 				SimpleSlot slot = getInstance(
-						new SimpleInstanceGateway(
+						new SimpleActorGateway(
 								TestingUtils.defaultExecutionContext())
 				).allocateSimpleSlot(jid);
 				ee.deployToSlot(slot);

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexCancelTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexCancelTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexCancelTest.java
index 9db330b..e9b67af 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexCancelTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexCancelTest.java
@@ -26,9 +26,9 @@ import java.io.IOException;
 
 import org.apache.flink.runtime.akka.AkkaUtils;
 import org.apache.flink.runtime.execution.ExecutionState;
-import org.apache.flink.runtime.instance.BaseTestingInstanceGateway;
-import org.apache.flink.runtime.instance.DummyInstanceGateway;
-import org.apache.flink.runtime.instance.InstanceGateway;
+import org.apache.flink.runtime.instance.BaseTestingActorGateway;
+import org.apache.flink.runtime.instance.DummyActorGateway;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.instance.SimpleSlot;
 import org.apache.flink.runtime.instance.Instance;
 import org.apache.flink.api.common.JobID;
@@ -125,12 +125,12 @@ public class ExecutionVertexCancelTest {
 			setVertexState(vertex, ExecutionState.SCHEDULED);
 			assertEquals(ExecutionState.SCHEDULED, vertex.getExecutionState());
 
-			InstanceGateway instanceGateway = new CancelSequenceInstanceGateway(
+			ActorGateway actorGateway = new CancelSequenceActorGateway(
 					executionContext,
 					new TaskOperationResult(execId, true),
 					new TaskOperationResult(execId, false));
 
-			Instance instance = getInstance(instanceGateway);
+			Instance instance = getInstance(actorGateway);
 			SimpleSlot slot = instance.allocateSimpleSlot(new JobID());
 
 			vertex.deployToSlot(slot);
@@ -195,13 +195,13 @@ public class ExecutionVertexCancelTest {
 
 			// task manager cancel sequence mock actor
 			// first return NOT SUCCESS (task not found, cancel call overtook deploy call), then success (cancel call after deploy call)
-			InstanceGateway instanceGateway = new CancelSequenceInstanceGateway(
+			ActorGateway actorGateway = new CancelSequenceActorGateway(
 					executionContext,
 					new	TaskOperationResult(execId, false),
 					new TaskOperationResult(execId, true)
 			);
 
-			Instance instance = getInstance(instanceGateway);
+			Instance instance = getInstance(actorGateway);
 			SimpleSlot slot = instance.allocateSimpleSlot(new JobID());
 
 			vertex.deployToSlot(slot);
@@ -258,11 +258,11 @@ public class ExecutionVertexCancelTest {
 					AkkaUtils.getDefaultTimeout());
 			final ExecutionAttemptID execId = vertex.getCurrentExecutionAttempt().getAttemptId();
 
-			InstanceGateway instanceGateway = new CancelSequenceInstanceGateway(
+			ActorGateway actorGateway = new CancelSequenceActorGateway(
 					TestingUtils.directExecutionContext(),
 					new TaskOperationResult(execId, true));
 
-			Instance instance = getInstance(instanceGateway);
+			Instance instance = getInstance(actorGateway);
 			SimpleSlot slot = instance.allocateSimpleSlot(new JobID());
 
 			setVertexState(vertex, ExecutionState.RUNNING);
@@ -299,12 +299,12 @@ public class ExecutionVertexCancelTest {
 					AkkaUtils.getDefaultTimeout());
 			final ExecutionAttemptID execId = vertex.getCurrentExecutionAttempt().getAttemptId();
 
-			final InstanceGateway instanceGateway = new CancelSequenceInstanceGateway(
+			final ActorGateway actorGateway = new CancelSequenceActorGateway(
 					TestingUtils.directExecutionContext(),
 					new TaskOperationResult(execId, true)
 			);
 
-			Instance instance = getInstance(instanceGateway);
+			Instance instance = getInstance(actorGateway);
 			SimpleSlot slot = instance.allocateSimpleSlot(new JobID());
 
 			setVertexState(vertex, ExecutionState.RUNNING);
@@ -350,12 +350,12 @@ public class ExecutionVertexCancelTest {
 			final ExecutionAttemptID execId = vertex.getCurrentExecutionAttempt().getAttemptId();
 
 
-			final InstanceGateway instanceGateway = new CancelSequenceInstanceGateway(
+			final ActorGateway actorGateway = new CancelSequenceActorGateway(
 					TestingUtils.directExecutionContext(),
 					new TaskOperationResult(execId, false)
 			);
 
-			Instance instance = getInstance(instanceGateway);
+			Instance instance = getInstance(actorGateway);
 			SimpleSlot slot = instance.allocateSimpleSlot(new JobID());
 
 			setVertexState(vertex, ExecutionState.RUNNING);
@@ -386,7 +386,7 @@ public class ExecutionVertexCancelTest {
 			final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
 					AkkaUtils.getDefaultTimeout());
 
-			final InstanceGateway gateway = new CancelSequenceInstanceGateway(TestingUtils.directExecutionContext());
+			final ActorGateway gateway = new CancelSequenceActorGateway(TestingUtils.directExecutionContext());
 
 			Instance instance = getInstance(gateway);
 			SimpleSlot slot = instance.allocateSimpleSlot(new JobID());
@@ -423,7 +423,7 @@ public class ExecutionVertexCancelTest {
 					AkkaUtils.getDefaultTimeout());
 			final ExecutionAttemptID execID = vertex.getCurrentExecutionAttempt().getAttemptId();
 
-			final InstanceGateway gateway = new CancelSequenceInstanceGateway(
+			final ActorGateway gateway = new CancelSequenceActorGateway(
 					TestingUtils.defaultExecutionContext(),
 					new TaskOperationResult(execID, true));
 
@@ -482,7 +482,7 @@ public class ExecutionVertexCancelTest {
 			// deploying after canceling from CREATED needs to raise an exception, because
 			// the scheduler (or any caller) needs to know that the slot should be released
 			try {
-				Instance instance = getInstance(DummyInstanceGateway.INSTANCE);
+				Instance instance = getInstance(DummyActorGateway.INSTANCE);
 				SimpleSlot slot = instance.allocateSimpleSlot(new JobID());
 
 				vertex.deployToSlot(slot);
@@ -525,7 +525,7 @@ public class ExecutionVertexCancelTest {
 						AkkaUtils.getDefaultTimeout());
 				setVertexState(vertex, ExecutionState.CANCELING);
 
-				Instance instance = getInstance(DummyInstanceGateway.INSTANCE);
+				Instance instance = getInstance(DummyActorGateway.INSTANCE);
 				SimpleSlot slot = instance.allocateSimpleSlot(new JobID());
 
 				vertex.deployToSlot(slot);
@@ -541,7 +541,7 @@ public class ExecutionVertexCancelTest {
 				ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
 						AkkaUtils.getDefaultTimeout());
 
-				Instance instance = getInstance(DummyInstanceGateway.INSTANCE);
+				Instance instance = getInstance(DummyActorGateway.INSTANCE);
 				SimpleSlot slot = instance.allocateSimpleSlot(new JobID());
 
 				setVertexResource(vertex, slot);
@@ -562,11 +562,11 @@ public class ExecutionVertexCancelTest {
 		}
 	}
 
-	public static class CancelSequenceInstanceGateway extends BaseTestingInstanceGateway {
+	public static class CancelSequenceActorGateway extends BaseTestingActorGateway {
 		private final TaskOperationResult[] results;
 		private int index = -1;
 
-		public CancelSequenceInstanceGateway(ExecutionContext executionContext, TaskOperationResult ... result) {
+		public CancelSequenceActorGateway(ExecutionContext executionContext, TaskOperationResult... result) {
 			super(executionContext);
 			this.results = result;
 		}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexDeploymentTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexDeploymentTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexDeploymentTest.java
index 431c3a9..81ec6c9 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexDeploymentTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexDeploymentTest.java
@@ -43,7 +43,7 @@ public class ExecutionVertexDeploymentTest {
 
 			// mock taskmanager to simply accept the call
 			Instance instance = getInstance(
-					new SimpleInstanceGateway(TestingUtils.directExecutionContext()));
+					new SimpleActorGateway(TestingUtils.directExecutionContext()));
 			final SimpleSlot slot = instance.allocateSimpleSlot(ejv.getJobId());
 
 			final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
@@ -81,7 +81,7 @@ public class ExecutionVertexDeploymentTest {
 			final ExecutionJobVertex ejv = getExecutionVertex(jid, TestingUtils.directExecutionContext());
 
 			final Instance instance = getInstance(
-					new SimpleInstanceGateway(TestingUtils.directExecutionContext()));
+					new SimpleActorGateway(TestingUtils.directExecutionContext()));
 			final SimpleSlot slot = instance.allocateSimpleSlot(ejv.getJobId());
 
 			final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
@@ -124,7 +124,7 @@ public class ExecutionVertexDeploymentTest {
 					AkkaUtils.getDefaultTimeout());
 
 			final Instance instance = getInstance(
-					new SimpleInstanceGateway(TestingUtils.defaultExecutionContext()));
+					new SimpleActorGateway(TestingUtils.defaultExecutionContext()));
 			final SimpleSlot slot = instance.allocateSimpleSlot(ejv.getJobId());
 
 			assertEquals(ExecutionState.CREATED, vertex.getExecutionState());
@@ -171,7 +171,7 @@ public class ExecutionVertexDeploymentTest {
 					AkkaUtils.getDefaultTimeout());
 
 			final Instance instance = getInstance(
-					new SimpleFailingInstanceGateway(TestingUtils.directExecutionContext()));
+					new SimpleFailingActorGateway(TestingUtils.directExecutionContext()));
 			final SimpleSlot slot = instance.allocateSimpleSlot(ejv.getJobId());
 
 			assertEquals(ExecutionState.CREATED, vertex.getExecutionState());
@@ -201,7 +201,7 @@ public class ExecutionVertexDeploymentTest {
 					AkkaUtils.getDefaultTimeout());
 
 			final Instance instance = getInstance(
-					new SimpleFailingInstanceGateway(TestingUtils.directExecutionContext()));
+					new SimpleFailingActorGateway(TestingUtils.directExecutionContext()));
 			final SimpleSlot slot = instance.allocateSimpleSlot(ejv.getJobId());
 
 			assertEquals(ExecutionState.CREATED, vertex.getExecutionState());
@@ -244,7 +244,7 @@ public class ExecutionVertexDeploymentTest {
 			final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
 					AkkaUtils.getDefaultTimeout());
 
-			final Instance instance = getInstance(new SimpleInstanceGateway(TestingUtils.directExecutionContext()));
+			final Instance instance = getInstance(new SimpleActorGateway(TestingUtils.directExecutionContext()));
 			final SimpleSlot slot = instance.allocateSimpleSlot(ejv.getJobId());
 
 			assertEquals(ExecutionState.CREATED, vertex.getExecutionState());
@@ -286,7 +286,7 @@ public class ExecutionVertexDeploymentTest {
 			final ExecutionAttemptID eid = vertex.getCurrentExecutionAttempt().getAttemptId();
 
 			final Instance instance = getInstance(
-					new ExecutionVertexCancelTest.CancelSequenceInstanceGateway(
+					new ExecutionVertexCancelTest.CancelSequenceActorGateway(
 							context,
 							new TaskOperationResult(eid, false),
 							new TaskOperationResult(eid, true)));

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexSchedulingTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexSchedulingTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexSchedulingTest.java
index 8ea7017..5e9ee33 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexSchedulingTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionVertexSchedulingTest.java
@@ -25,7 +25,7 @@ import static org.mockito.Mockito.*;
 
 import org.apache.flink.runtime.akka.AkkaUtils;
 import org.apache.flink.runtime.execution.ExecutionState;
-import org.apache.flink.runtime.instance.DummyInstanceGateway;
+import org.apache.flink.runtime.instance.DummyActorGateway;
 import org.apache.flink.runtime.instance.Instance;
 import org.apache.flink.runtime.instance.SimpleSlot;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
@@ -47,7 +47,7 @@ public class ExecutionVertexSchedulingTest {
 					AkkaUtils.getDefaultTimeout());
 
 			// a slot than cannot be deployed to
-			final Instance instance = getInstance(DummyInstanceGateway.INSTANCE);
+			final Instance instance = getInstance(DummyActorGateway.INSTANCE);
 			final SimpleSlot slot = instance.allocateSimpleSlot(ejv.getJobId());
 			
 			slot.releaseSlot();
@@ -77,7 +77,7 @@ public class ExecutionVertexSchedulingTest {
 					AkkaUtils.getDefaultTimeout());
 
 			// a slot than cannot be deployed to
-			final Instance instance = getInstance(DummyInstanceGateway.INSTANCE);
+			final Instance instance = getInstance(DummyActorGateway.INSTANCE);
 			final SimpleSlot slot = instance.allocateSimpleSlot(ejv.getJobId());
 
 			slot.releaseSlot();
@@ -113,7 +113,7 @@ public class ExecutionVertexSchedulingTest {
 			final ExecutionVertex vertex = new ExecutionVertex(ejv, 0, new IntermediateResult[0],
 					AkkaUtils.getDefaultTimeout());
 
-			final Instance instance = getInstance(new ExecutionGraphTestUtils.SimpleInstanceGateway(TestingUtils.defaultExecutionContext()));
+			final Instance instance = getInstance(new ExecutionGraphTestUtils.SimpleActorGateway(TestingUtils.defaultExecutionContext()));
 			final SimpleSlot slot = instance.allocateSimpleSlot(ejv.getJobId());
 
 			Scheduler scheduler = mock(Scheduler.class);

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/LocalInputSplitsTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/LocalInputSplitsTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/LocalInputSplitsTest.java
index b4a7e63..2530a53 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/LocalInputSplitsTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/LocalInputSplitsTest.java
@@ -377,7 +377,7 @@ public class LocalInputSplitsTest {
 		when(connection.getFQDNHostname()).thenReturn(hostname);
 		
 		return new Instance(
-				new ExecutionGraphTestUtils.SimpleInstanceGateway(
+				new ExecutionGraphTestUtils.SimpleActorGateway(
 						TestingUtils.defaultExecutionContext()),
 				connection,
 				new InstanceID(),

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/TerminalStateDeadlockTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/TerminalStateDeadlockTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/TerminalStateDeadlockTest.java
index b779d79..f42543f 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/TerminalStateDeadlockTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/TerminalStateDeadlockTest.java
@@ -21,7 +21,7 @@ package org.apache.flink.runtime.executiongraph;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.execution.ExecutionState;
-import org.apache.flink.runtime.instance.DummyInstanceGateway;
+import org.apache.flink.runtime.instance.DummyActorGateway;
 import org.apache.flink.runtime.instance.HardwareDescription;
 import org.apache.flink.runtime.instance.Instance;
 import org.apache.flink.runtime.instance.InstanceConnectionInfo;
@@ -78,7 +78,7 @@ public class TerminalStateDeadlockTest {
 			InstanceConnectionInfo ci = new InstanceConnectionInfo(address, 12345);
 				
 			HardwareDescription resources = new HardwareDescription(4, 4000000, 3000000, 2000000);
-			Instance instance = new Instance(DummyInstanceGateway.INSTANCE, ci, new InstanceID(), resources, 4);
+			Instance instance = new Instance(DummyActorGateway.INSTANCE, ci, new InstanceID(), resources, 4);
 
 			this.resource = instance.allocateSimpleSlot(new JobID());
 		}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/VertexLocationConstraintTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/VertexLocationConstraintTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/VertexLocationConstraintTest.java
index 3305254..8604b63 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/VertexLocationConstraintTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/VertexLocationConstraintTest.java
@@ -26,7 +26,7 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.concurrent.TimeUnit;
 
-import org.apache.flink.runtime.instance.DummyInstanceGateway;
+import org.apache.flink.runtime.instance.DummyActorGateway;
 import org.apache.flink.runtime.instance.HardwareDescription;
 import org.apache.flink.runtime.instance.Instance;
 import org.apache.flink.runtime.instance.InstanceConnectionInfo;
@@ -403,7 +403,7 @@ public class VertexLocationConstraintTest {
 			
 			ExecutionVertex ev = eg.getAllVertices().get(vertex.getID()).getTaskVertices()[0];
 			
-			Instance instance = ExecutionGraphTestUtils.getInstance(DummyInstanceGateway.INSTANCE);
+			Instance instance = ExecutionGraphTestUtils.getInstance(DummyActorGateway.INSTANCE);
 			ev.setLocationConstraintHosts(Collections.singletonList(instance));
 			
 			assertNotNull(ev.getPreferredLocations());
@@ -435,7 +435,7 @@ public class VertexLocationConstraintTest {
 		when(connection.getFQDNHostname()).thenReturn(hostname);
 		
 		return new Instance(
-				new ExecutionGraphTestUtils.SimpleInstanceGateway(
+				new ExecutionGraphTestUtils.SimpleActorGateway(
 						TestingUtils.defaultExecutionContext()),
 				connection,
 				new InstanceID(),

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/instance/BaseTestingActorGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/BaseTestingActorGateway.java b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/BaseTestingActorGateway.java
new file mode 100644
index 0000000..2e62781
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/BaseTestingActorGateway.java
@@ -0,0 +1,120 @@
+/*
+ * 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.instance;
+
+import akka.actor.ActorRef;
+import akka.dispatch.Futures;
+import scala.Option;
+import scala.concurrent.ExecutionContext;
+import scala.concurrent.Future;
+import scala.concurrent.duration.FiniteDuration;
+
+import java.util.UUID;
+import java.util.concurrent.Callable;
+
+/**
+ * Abstract base class for testing {@link ActorGateway} instances. The implementing subclass
+ * only has to provide an implementation for handleMessage which contains the logic to treat
+ * different messages.
+ */
+abstract public class BaseTestingActorGateway implements ActorGateway {
+	/**
+	 * {@link ExecutionContext} which is used to execute the futures.
+	 */
+	private final ExecutionContext executionContext;
+
+	public BaseTestingActorGateway(ExecutionContext executionContext) {
+		this.executionContext = executionContext;
+	}
+
+	@Override
+	public Future<Object> ask(Object message, FiniteDuration timeout) {
+		try {
+			final Object result = handleMessage(message);
+
+			return Futures.future(new Callable<Object>() {
+				@Override
+				public Object call() throws Exception {
+					return result;
+				}
+			}, executionContext);
+
+		} catch (final Exception e) {
+			// if an exception occurred in the handleMessage method then return it as part of the future
+			return Futures.future(new Callable<Object>() {
+				@Override
+				public Object call() throws Exception {
+					throw e;
+				}
+			}, executionContext);
+		}
+	}
+
+	/**
+	 * Handles the supported messages by this InstanceGateway
+	 *
+	 * @param message Message to handle
+	 * @return Result
+	 * @throws Exception
+	 */
+	abstract public Object handleMessage(Object message) throws Exception;
+
+	@Override
+	public void tell(Object message) {
+		try {
+			handleMessage(message);
+		} catch (Exception e) {
+			// discard exception because it happens on the "remote" instance
+		}
+	}
+
+	@Override
+	public void tell(Object message, ActorGateway sender) {
+		try{
+			handleMessage(message);
+		} catch (Exception e) {
+			// discard exception because it happens on the "remote" instance
+		}
+	}
+
+	@Override
+	public void forward(Object message, ActorGateway sender) {
+		throw new UnsupportedOperationException();
+	}
+
+	@Override
+	public Future<Object> retry(Object message, int numberRetries, FiniteDuration timeout, ExecutionContext executionContext) {
+		return ask(message, timeout);
+	}
+
+	@Override
+	public String path() {
+		return "BaseTestingInstanceGateway";
+	}
+
+	@Override
+	public ActorRef actor() {
+		return ActorRef.noSender();
+	}
+
+	@Override
+	public Option<UUID> leaderSessionID() {
+		return Option.empty();
+	}
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/instance/BaseTestingInstanceGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/BaseTestingInstanceGateway.java b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/BaseTestingInstanceGateway.java
deleted file mode 100644
index e9f8259..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/BaseTestingInstanceGateway.java
+++ /dev/null
@@ -1,94 +0,0 @@
-/*
- * 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.instance;
-
-import akka.actor.ActorPath;
-import akka.actor.ActorRef;
-import akka.dispatch.Futures;
-import scala.concurrent.ExecutionContext;
-import scala.concurrent.Future;
-import scala.concurrent.duration.FiniteDuration;
-
-import java.util.concurrent.Callable;
-
-/**
- * Abstract base class for testing {@link InstanceGateway} instances. The implementing subclass
- * only has to provide an implementation for handleMessage which contains the logic to treat
- * different messages.
- */
-abstract public class BaseTestingInstanceGateway implements InstanceGateway {
-	/**
-	 * {@link ExecutionContext} which is used to execute the futures.
-	 */
-	private final ExecutionContext executionContext;
-
-	public BaseTestingInstanceGateway(ExecutionContext executionContext) {
-		this.executionContext = executionContext;
-	}
-
-	@Override
-	public Future<Object> ask(Object message, FiniteDuration timeout) {
-		try {
-			final Object result = handleMessage(message);
-
-			return Futures.future(new Callable<Object>() {
-				@Override
-				public Object call() throws Exception {
-					return result;
-				}
-			}, executionContext);
-
-		} catch (final Exception e) {
-			// if an exception occurred in the handleMessage method then return it as part of the future
-			return Futures.future(new Callable<Object>() {
-				@Override
-				public Object call() throws Exception {
-					throw e;
-				}
-			}, executionContext);
-		}
-	}
-
-	/**
-	 * Handles the supported messages by this InstanceGateway
-	 *
-	 * @param message Message to handle
-	 * @return Result
-	 * @throws Exception
-	 */
-	abstract public Object handleMessage(Object message) throws Exception;
-
-	@Override
-	public void tell(Object message) {}
-
-	@Override
-	public void forward(Object message, ActorRef sender) {
-		throw new UnsupportedOperationException();
-	}
-
-	@Override
-	public Future<Object> retry(Object message, int numberRetries, FiniteDuration timeout, ExecutionContext executionContext) {
-		return ask(message, timeout);
-	}
-
-	@Override
-	public String path() {
-		throw new UnsupportedOperationException();
-	}
-}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/instance/DummyActorGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/DummyActorGateway.java b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/DummyActorGateway.java
new file mode 100644
index 0000000..10762f2
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/DummyActorGateway.java
@@ -0,0 +1,68 @@
+/*
+ * 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.instance;
+
+import akka.actor.ActorRef;
+import scala.Option;
+import scala.concurrent.ExecutionContext;
+import scala.concurrent.Future;
+import scala.concurrent.duration.FiniteDuration;
+
+import java.util.UUID;
+
+/**
+ * Dummy {@link ActorGateway} implementation used for testing.
+ */
+public class DummyActorGateway implements ActorGateway {
+	public static final DummyActorGateway INSTANCE = new DummyActorGateway();
+
+	@Override
+	public Future<Object> ask(Object message, FiniteDuration timeout) {
+		return null;
+	}
+
+	@Override
+	public void tell(Object message) {}
+
+	@Override
+	public void tell(Object message, ActorGateway sender) {}
+
+	@Override
+	public void forward(Object message, ActorGateway sender) {}
+
+	@Override
+	public Future<Object> retry(Object message, int numberRetries, FiniteDuration timeout, ExecutionContext executionContext) {
+		return null;
+	}
+
+	@Override
+	public String path() {
+		return "DummyInstanceGateway";
+	}
+
+	@Override
+	public ActorRef actor() {
+		return ActorRef.noSender();
+	}
+
+	@Override
+	public Option<UUID> leaderSessionID() {
+		return Option.<UUID>empty();
+	}
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/instance/DummyInstanceGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/DummyInstanceGateway.java b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/DummyInstanceGateway.java
deleted file mode 100644
index 5941201..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/DummyInstanceGateway.java
+++ /dev/null
@@ -1,57 +0,0 @@
-/*
- * 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.instance;
-
-import akka.actor.ActorPath;
-import akka.actor.ActorRef;
-import scala.concurrent.ExecutionContext;
-import scala.concurrent.Future;
-import scala.concurrent.duration.FiniteDuration;
-
-/**
- * Dummy {@link InstanceGateway} implementation used for testing.
- */
-public class DummyInstanceGateway implements InstanceGateway {
-	public static final DummyInstanceGateway INSTANCE = new DummyInstanceGateway();
-
-	@Override
-	public Future<Object> ask(Object message, FiniteDuration timeout) {
-		throw new UnsupportedOperationException();
-	}
-
-	@Override
-	public void tell(Object message) {
-		throw new UnsupportedOperationException();
-	}
-
-	@Override
-	public void forward(Object message, ActorRef sender) {
-		throw new UnsupportedOperationException();
-	}
-
-	@Override
-	public Future<Object> retry(Object message, int numberRetries, FiniteDuration timeout, ExecutionContext executionContext) {
-		throw new UnsupportedOperationException();
-	}
-
-	@Override
-	public String path() {
-		return "DummyInstanceGateway";
-	}
-}


[3/6] flink git commit: [FLINK-2332] [runtime] Adds leader session IDs and registration session IDs to JobManager and TaskManager messages.

Posted by tr...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/instance/InstanceManagerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/InstanceManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/InstanceManagerTest.java
index a0d0a83..2c01f1f 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/InstanceManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/InstanceManagerTest.java
@@ -31,6 +31,7 @@ import java.util.Collection;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Set;
+import java.util.UUID;
 
 import akka.actor.ActorSystem;
 import akka.testkit.JavaTestKit;
@@ -41,6 +42,8 @@ import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import scala.Option;
+import scala.Some;
 
 /**
  * Tests for {@link org.apache.flink.runtime.instance.InstanceManager}.
@@ -49,6 +52,8 @@ public class InstanceManagerTest{
 
 	static ActorSystem system;
 
+	static Option<UUID> leaderSessionID = new Some<UUID>(UUID.randomUUID());
+
 	@BeforeClass
 	public static void setup(){
 		system = ActorSystem.create("TestingActorSystem", TestingUtils.testConfig());
@@ -80,9 +85,9 @@ public class InstanceManagerTest{
 			final JavaTestKit probe2 = new JavaTestKit(system);
 			final JavaTestKit probe3 = new JavaTestKit(system);
 
-			cm.registerTaskManager(probe1.getRef(), ici1, hardwareDescription, 1);
-			cm.registerTaskManager(probe2.getRef(), ici2, hardwareDescription, 2);
-			cm.registerTaskManager(probe3.getRef(), ici3, hardwareDescription, 5);
+			cm.registerTaskManager(probe1.getRef(), ici1, hardwareDescription, 1, leaderSessionID);
+			cm.registerTaskManager(probe2.getRef(), ici2, hardwareDescription, 2, leaderSessionID);
+			cm.registerTaskManager(probe3.getRef(), ici3, hardwareDescription, 5, leaderSessionID);
 			
 			assertEquals(3, cm.getNumberOfRegisteredTaskManagers());
 			assertEquals(8, cm.getTotalNumberOfSlots());
@@ -120,13 +125,13 @@ public class InstanceManagerTest{
 			InstanceConnectionInfo ici = new InstanceConnectionInfo(address, dataPort);
 
 			JavaTestKit probe = new JavaTestKit(system);
-			InstanceID i = cm.registerTaskManager(probe.getRef(), ici, resources, 1);
+			InstanceID i = cm.registerTaskManager(probe.getRef(), ici, resources, 1, leaderSessionID);
 
 			assertNotNull(i);
 			assertEquals(1, cm.getNumberOfRegisteredTaskManagers());
 			assertEquals(1, cm.getTotalNumberOfSlots());
 			
-			InstanceID next = cm.registerTaskManager(probe.getRef(), ici, resources, 1);
+			InstanceID next = cm.registerTaskManager(probe.getRef(), ici, resources, 1, leaderSessionID);
 			assertNull(next);
 			
 			assertEquals(1, cm.getNumberOfRegisteredTaskManagers());
@@ -161,9 +166,9 @@ public class InstanceManagerTest{
 			JavaTestKit probe2 = new JavaTestKit(system);
 			JavaTestKit probe3 = new JavaTestKit(system);
 			
-			InstanceID i1 = cm.registerTaskManager(probe1.getRef(), ici1, hardwareDescription, 1);
-			InstanceID i2 = cm.registerTaskManager(probe2.getRef(), ici2, hardwareDescription, 1);
-			InstanceID i3 = cm.registerTaskManager(probe3.getRef(), ici3, hardwareDescription, 1);
+			InstanceID i1 = cm.registerTaskManager(probe1.getRef(), ici1, hardwareDescription, 1, leaderSessionID);
+			InstanceID i2 = cm.registerTaskManager(probe2.getRef(), ici2, hardwareDescription, 1, leaderSessionID);
+			InstanceID i3 = cm.registerTaskManager(probe3.getRef(), ici3, hardwareDescription, 1, leaderSessionID);
 
 			// report some immediate heart beats
 			assertTrue(cm.reportHeartBeat(i1, new byte[] {}));
@@ -216,7 +221,7 @@ public class InstanceManagerTest{
 				InstanceConnectionInfo ici = new InstanceConnectionInfo(address, 20000);
 
 				JavaTestKit probe = new JavaTestKit(system);
-				cm.registerTaskManager(probe.getRef(), ici, resources, 1);
+				cm.registerTaskManager(probe.getRef(), ici, resources, 1, leaderSessionID);
 				fail("Should raise exception in shutdown state");
 			}
 			catch (IllegalStateException e) {

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/instance/InstanceTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/InstanceTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/InstanceTest.java
index c075a17..2c0f4d9 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/InstanceTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/InstanceTest.java
@@ -38,7 +38,7 @@ public class InstanceTest {
 			InetAddress address = InetAddress.getByName("127.0.0.1");
 			InstanceConnectionInfo connection = new InstanceConnectionInfo(address, 10001);
 
-			Instance instance = new Instance(DummyInstanceGateway.INSTANCE, connection, new InstanceID(), hardwareDescription, 4);
+			Instance instance = new Instance(DummyActorGateway.INSTANCE, connection, new InstanceID(), hardwareDescription, 4);
 
 			assertEquals(4, instance.getTotalNumberOfSlots());
 			assertEquals(4, instance.getNumberOfAvailableSlots());
@@ -99,7 +99,7 @@ public class InstanceTest {
 			InetAddress address = InetAddress.getByName("127.0.0.1");
 			InstanceConnectionInfo connection = new InstanceConnectionInfo(address, 10001);
 
-			Instance instance = new Instance(DummyInstanceGateway.INSTANCE, connection, new InstanceID(), hardwareDescription, 3);
+			Instance instance = new Instance(DummyActorGateway.INSTANCE, connection, new InstanceID(), hardwareDescription, 3);
 
 			assertEquals(3, instance.getNumberOfAvailableSlots());
 
@@ -129,7 +129,7 @@ public class InstanceTest {
 			InetAddress address = InetAddress.getByName("127.0.0.1");
 			InstanceConnectionInfo connection = new InstanceConnectionInfo(address, 10001);
 
-			Instance instance = new Instance(DummyInstanceGateway.INSTANCE, connection, new InstanceID(), hardwareDescription, 3);
+			Instance instance = new Instance(DummyActorGateway.INSTANCE, connection, new InstanceID(), hardwareDescription, 3);
 
 			assertEquals(3, instance.getNumberOfAvailableSlots());
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/instance/SimpleSlotTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/SimpleSlotTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/SimpleSlotTest.java
index 29ec7b5..997bb46 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/instance/SimpleSlotTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/instance/SimpleSlotTest.java
@@ -146,7 +146,7 @@ public class SimpleSlotTest {
 		InetAddress address = InetAddress.getByName("127.0.0.1");
 		InstanceConnectionInfo connection = new InstanceConnectionInfo(address, 10001);
 
-		Instance instance = new Instance(DummyInstanceGateway.INSTANCE, connection, new InstanceID(), hardwareDescription, 1);
+		Instance instance = new Instance(DummyActorGateway.INSTANCE, connection, new InstanceID(), hardwareDescription, 1);
 		return instance.allocateSimpleSlot(new JobID());
 	}
 }
\ No newline at end of file

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentTest.java
index 6affdcc..d657ebf 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/NetworkEnvironmentTest.java
@@ -20,8 +20,8 @@ package org.apache.flink.runtime.io.network;
 
 import static org.junit.Assert.*;
 
-import akka.actor.ActorRef;
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.runtime.instance.DummyActorGateway;
 import org.apache.flink.runtime.io.disk.iomanager.IOManager;
 import org.apache.flink.runtime.io.network.buffer.BufferPool;
 import org.apache.flink.runtime.io.network.netty.NettyConfig;
@@ -29,7 +29,6 @@ import org.apache.flink.runtime.net.NetUtils;
 import org.apache.flink.runtime.taskmanager.NetworkEnvironmentConfiguration;
 import org.apache.flink.runtime.testingUtils.TestingUtils;
 import org.junit.Test;
-import org.mockito.Mockito;
 import scala.Some;
 import scala.Tuple2;
 import scala.concurrent.duration.FiniteDuration;
@@ -78,9 +77,9 @@ public class NetworkEnvironmentTest {
 			assertNull(env.getPartitionManager());
 
 			// associate the environment with some mock actors
-			ActorRef jmActor = Mockito.mock(ActorRef.class);
-			ActorRef tmActor = Mockito.mock(ActorRef.class);
-			env.associateWithTaskManagerAndJobManager(jmActor, tmActor);
+			env.associateWithTaskManagerAndJobManager(
+					DummyActorGateway.INSTANCE,
+					DummyActorGateway.INSTANCE);
 
 			assertNotNull(env.getConnectionManager());
 			assertNotNull(env.getPartitionConsumableNotifier());
@@ -103,9 +102,10 @@ public class NetworkEnvironmentTest {
 			assertTrue(localPool.isDestroyed());
 
 			// associate once again
-			jmActor = Mockito.mock(ActorRef.class);
-			tmActor = Mockito.mock(ActorRef.class);
-			env.associateWithTaskManagerAndJobManager(jmActor, tmActor);
+			env.associateWithTaskManagerAndJobManager(
+					DummyActorGateway.INSTANCE,
+					DummyActorGateway.INSTANCE
+			);
 
 			assertNotNull(env.getConnectionManager());
 			assertNotNull(env.getPartitionConsumableNotifier());

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartialConsumePipelinedResultTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartialConsumePipelinedResultTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartialConsumePipelinedResultTest.java
index 7e15af4..cf81e3e 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartialConsumePipelinedResultTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartialConsumePipelinedResultTest.java
@@ -102,7 +102,7 @@ public class PartialConsumePipelinedResultTest {
 
 		JobClient.submitJobAndWait(
 				jobClient,
-				flink.getJobManager(),
+				flink.getJobManagerGateway(),
 				jobGraph,
 				TestingUtils.TESTING_DURATION(),
 				false);

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobManagerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobManagerTest.java
index 3eb4be8..8b01f06 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobManagerTest.java
@@ -18,7 +18,6 @@
 
 package org.apache.flink.runtime.jobmanager;
 
-import akka.actor.ActorRef;
 import akka.actor.ActorSystem;
 import akka.actor.Status;
 import akka.testkit.JavaTestKit;
@@ -32,11 +31,14 @@ import org.apache.flink.runtime.executiongraph.ExecutionGraph;
 import org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils;
 import org.apache.flink.runtime.executiongraph.ExecutionVertex;
 import org.apache.flink.runtime.executiongraph.IntermediateResultPartition;
+import org.apache.flink.runtime.instance.AkkaActorGateway;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
 import org.apache.flink.runtime.jobgraph.JobVertex;
 import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
 import org.apache.flink.runtime.jobgraph.JobGraph;
-import org.apache.flink.runtime.messages.JobManagerMessages;
+import org.apache.flink.runtime.messages.JobManagerMessages.LeaderSessionMessage;
+import org.apache.flink.runtime.messages.JobManagerMessages.SubmitJob;
 import org.apache.flink.runtime.messages.JobManagerMessages.RequestPartitionState;
 import org.apache.flink.runtime.messages.TaskMessages.PartitionState;
 import org.apache.flink.runtime.testingUtils.TestingCluster;
@@ -47,10 +49,12 @@ import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages.WaitForAl
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import scala.Option;
 import scala.Some;
 import scala.Tuple2;
 
 import java.net.InetAddress;
+import java.util.UUID;
 
 import static org.apache.flink.runtime.io.network.partition.ResultPartitionType.PIPELINED;
 import static org.apache.flink.runtime.testingUtils.TestingUtils.DEFAULT_AKKA_ASK_TIMEOUT;
@@ -113,14 +117,18 @@ public class JobManagerTest {
 				final JobGraph jobGraph = new JobGraph("Blocking test job", sender);
 				final JobID jid = jobGraph.getJobID();
 
-				final ActorRef jm = cluster.getJobManager();
+				final ActorGateway jobManagerGateway = cluster.getJobManagerGateway();
+
+				// we can set the leader session ID to None because we don't use this gateway to send messages
+				final ActorGateway testActorGateway = new AkkaActorGateway(getTestActor(), Option.<UUID>empty());
 
 				// Submit the job and wait for all vertices to be running
-				jm.tell(new JobManagerMessages.SubmitJob(jobGraph, false), getTestActor());
+				jobManagerGateway.tell(new SubmitJob(jobGraph, false), testActorGateway);
 				expectMsgClass(Status.Success.class);
 
-				jm.tell(new WaitForAllVerticesToBeRunningOrFinished(jobGraph.getJobID()),
-						getTestActor());
+				jobManagerGateway.tell(
+						new WaitForAllVerticesToBeRunningOrFinished(jobGraph.getJobID()),
+						testActorGateway);
 
 				expectMsgClass(TestingJobManagerMessages.AllVerticesRunning.class);
 
@@ -128,7 +136,7 @@ public class JobManagerTest {
 				final ExecutionAttemptID receiver = new ExecutionAttemptID();
 
 				// Request the execution graph to get the runtime info
-				jm.tell(new RequestExecutionGraph(jid), getTestActor());
+				jobManagerGateway.tell(new RequestExecutionGraph(jid), testActorGateway);
 
 				final ExecutionGraph eg = expectMsgClass(ExecutionGraphFound.class)
 						.executionGraph();
@@ -152,9 +160,13 @@ public class JobManagerTest {
 				for (ExecutionState state : ExecutionState.values()) {
 					ExecutionGraphTestUtils.setVertexState(vertex, state);
 
-					jm.tell(request, getTestActor());
+					jobManagerGateway.tell(request, testActorGateway);
+
+					LeaderSessionMessage lsm = expectMsgClass(LeaderSessionMessage.class);
 
-					PartitionState resp = expectMsgClass(PartitionState.class);
+					assertEquals(PartitionState.class, lsm.message().getClass());
+
+					PartitionState resp = (PartitionState) lsm.message();
 
 					assertEquals(request.taskExecutionId(), resp.taskExecutionId());
 					assertEquals(request.taskResultId(), resp.taskResultId());
@@ -165,9 +177,13 @@ public class JobManagerTest {
 				// 2. Non-existing execution
 				request = new RequestPartitionState(jid, new ResultPartitionID(), receiver, rid);
 
-				jm.tell(request, getTestActor());
+				jobManagerGateway.tell(request, testActorGateway);
+
+				LeaderSessionMessage lsm = expectMsgClass(LeaderSessionMessage.class);
 
-				PartitionState resp = expectMsgClass(PartitionState.class);
+				assertEquals(PartitionState.class, lsm.message().getClass());
+
+				PartitionState resp = (PartitionState) lsm.message();
 
 				assertEquals(request.taskExecutionId(), resp.taskExecutionId());
 				assertEquals(request.taskResultId(), resp.taskResultId());
@@ -178,9 +194,13 @@ public class JobManagerTest {
 				request = new RequestPartitionState(
 						new JobID(), new ResultPartitionID(), receiver, rid);
 
-				jm.tell(request, getTestActor());
+				jobManagerGateway.tell(request, testActorGateway);
+
+				lsm = expectMsgClass(LeaderSessionMessage.class);
+
+				assertEquals(PartitionState.class, lsm.message().getClass());
 
-				resp = expectMsgClass(PartitionState.class);
+				resp = (PartitionState) lsm.message();
 
 				assertEquals(request.taskExecutionId(), resp.taskExecutionId());
 				assertEquals(request.taskResultId(), resp.taskResultId());

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobSubmitTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobSubmitTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobSubmitTest.java
index 2f11c08..560fc30 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobSubmitTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobSubmitTest.java
@@ -20,13 +20,13 @@ package org.apache.flink.runtime.jobmanager;
 
 import akka.actor.ActorRef;
 import akka.actor.ActorSystem;
-import akka.pattern.Patterns;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.StreamingMode;
 import org.apache.flink.runtime.akka.AkkaUtils;
 import org.apache.flink.runtime.blob.BlobClient;
 import org.apache.flink.runtime.blob.BlobKey;
 import org.apache.flink.runtime.client.JobExecutionException;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.jobgraph.JobVertex;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.runtime.messages.JobManagerMessages;
@@ -51,10 +51,10 @@ import static org.junit.Assert.fail;
  */
 public class JobSubmitTest {
 
-	private static final long TIMEOUT = 5000;
+	private static final FiniteDuration timeout = new FiniteDuration(5000, TimeUnit.MILLISECONDS);
 
 	private static ActorSystem jobManagerSystem;
-	private static ActorRef jobManager;
+	private static ActorGateway jobManager;
 
 	@BeforeClass
 	public static void setupJobManager() {
@@ -62,7 +62,18 @@ public class JobSubmitTest {
 
 		scala.Option<Tuple2<String, Object>> listeningAddress = scala.Option.empty();
 		jobManagerSystem = AkkaUtils.createActorSystem(config, listeningAddress);
-		jobManager = JobManager.startJobManagerActors(config, jobManagerSystem, StreamingMode.BATCH_ONLY)._1();
+		ActorRef jobManagerActorRef = JobManager.startJobManagerActors(
+				config,
+				jobManagerSystem,
+				StreamingMode.BATCH_ONLY)._1();
+
+		try {
+			jobManager = JobManager.getJobManagerGateway(jobManagerActorRef, timeout);
+		} catch (Exception e) {
+			fail("Could not retrieve the JobManager gateway. " + e.getMessage());
+		}
+
+
 	}
 
 	@AfterClass
@@ -81,8 +92,8 @@ public class JobSubmitTest {
 			JobGraph jg = new JobGraph("test job", jobVertex);
 
 			// request the blob port from the job manager
-			Future<Object> future = Patterns.ask(jobManager, JobManagerMessages.getRequestBlobManagerPort(), TIMEOUT);
-			int blobPort = (Integer) Await.result(future, new FiniteDuration(TIMEOUT, TimeUnit.MILLISECONDS));
+			Future<Object> future = jobManager.ask(JobManagerMessages.getRequestBlobManagerPort(), timeout);
+			int blobPort = (Integer) Await.result(future, timeout);
 
 			// upload two dummy bytes and add their keys to the job graph as dependencies
 			BlobKey key1, key2;
@@ -102,10 +113,9 @@ public class JobSubmitTest {
 			jg.addBlob(key2);
 
 			// submit the job
-			Future<Object> submitFuture = Patterns.ask(jobManager,
-					new JobManagerMessages.SubmitJob(jg, false), TIMEOUT);
+			Future<Object> submitFuture = jobManager.ask(new JobManagerMessages.SubmitJob(jg, false), timeout);
 			try {
-				Await.result(submitFuture, new FiniteDuration(TIMEOUT, TimeUnit.MILLISECONDS));
+				Await.result(submitFuture, timeout);
 			}
 			catch (JobExecutionException e) {
 				// that is what we expect
@@ -142,10 +152,9 @@ public class JobSubmitTest {
 			JobGraph jg = new JobGraph("test job", jobVertex);
 
 			// submit the job
-			Future<Object> submitFuture = Patterns.ask(jobManager,
-					new JobManagerMessages.SubmitJob(jg, false), TIMEOUT);
+			Future<Object> submitFuture = jobManager.ask(new JobManagerMessages.SubmitJob(jg, false), timeout);
 			try {
-				Await.result(submitFuture, new FiniteDuration(TIMEOUT, TimeUnit.MILLISECONDS));
+				Await.result(submitFuture, timeout);
 			}
 			catch (JobExecutionException e) {
 				// that is what we expect

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/SchedulerTestUtils.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/SchedulerTestUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/SchedulerTestUtils.java
index 2de0635..b7bc860 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/SchedulerTestUtils.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/SchedulerTestUtils.java
@@ -31,7 +31,7 @@ import java.util.concurrent.atomic.AtomicInteger;
 
 import org.apache.flink.runtime.executiongraph.Execution;
 import org.apache.flink.runtime.executiongraph.ExecutionVertex;
-import org.apache.flink.runtime.instance.DummyInstanceGateway;
+import org.apache.flink.runtime.instance.DummyActorGateway;
 import org.apache.flink.runtime.instance.HardwareDescription;
 import org.apache.flink.runtime.instance.Instance;
 import org.apache.flink.runtime.instance.InstanceConnectionInfo;
@@ -66,7 +66,7 @@ public class SchedulerTestUtils {
 		final long GB = 1024L*1024*1024;
 		HardwareDescription resources = new HardwareDescription(4, 4*GB, 3*GB, 2*GB);
 		
-		return new Instance(DummyInstanceGateway.INSTANCE, ci, new InstanceID(), resources, numSlots);
+		return new Instance(DummyActorGateway.INSTANCE, ci, new InstanceID(), resources, numSlots);
 	}
 	
 	

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/ForwardingActorGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/ForwardingActorGateway.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/ForwardingActorGateway.java
new file mode 100644
index 0000000..e7e105f
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/ForwardingActorGateway.java
@@ -0,0 +1,47 @@
+/*
+ * 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.taskmanager;
+
+import org.apache.flink.runtime.instance.BaseTestingActorGateway;
+import org.apache.flink.runtime.testingUtils.TestingUtils;
+
+import java.util.concurrent.BlockingQueue;
+
+/**
+ * Testing {@link org.apache.flink.runtime.instance.ActorGateway} which stores all handled messages
+ * in a {@link BlockingQueue}.
+ */
+public class ForwardingActorGateway extends BaseTestingActorGateway {
+
+	private final BlockingQueue<Object> queue;
+
+	public ForwardingActorGateway(BlockingQueue<Object> queue) {
+		super(TestingUtils.directExecutionContext());
+
+		this.queue = queue;
+	}
+
+	@Override
+	public Object handleMessage(Object message) throws Exception {
+		// don't do anything with the message, but storing it in queue
+		queue.add(message);
+
+		return null;
+	}
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
index b18bf2a..10a33c3 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
@@ -18,13 +18,8 @@
 
 package org.apache.flink.runtime.taskmanager;
 
-import akka.actor.ActorSystem;
-import akka.actor.Props;
-import akka.actor.UntypedActor;
-
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.configuration.Configuration;
-import org.apache.flink.runtime.akka.AkkaUtils;
 import org.apache.flink.runtime.blob.BlobKey;
 import org.apache.flink.runtime.broadcast.BroadcastVariableManager;
 import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
@@ -34,6 +29,7 @@ import org.apache.flink.runtime.execution.ExecutionState;
 import org.apache.flink.runtime.execution.librarycache.LibraryCacheManager;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
 import org.apache.flink.runtime.filecache.FileCache;
+import org.apache.flink.runtime.instance.DummyActorGateway;
 import org.apache.flink.runtime.io.disk.iomanager.IOManager;
 import org.apache.flink.runtime.io.network.NetworkEnvironment;
 import org.apache.flink.runtime.io.network.partition.ResultPartitionConsumableNotifier;
@@ -44,9 +40,7 @@ import org.apache.flink.runtime.jobgraph.tasks.CheckpointNotificationOperator;
 import org.apache.flink.runtime.jobgraph.tasks.CheckpointedOperator;
 import org.apache.flink.runtime.memorymanager.MemoryManager;
 
-import org.junit.AfterClass;
 import org.junit.Before;
-import org.junit.BeforeClass;
 import org.junit.Test;
 
 import scala.concurrent.duration.FiniteDuration;
@@ -65,26 +59,9 @@ public class TaskAsyncCallTest {
 
 	private static final int NUM_CALLS = 1000;
 	
-	private static ActorSystem actorSystem;
-	
 	private static OneShotLatch awaitLatch;
 	private static OneShotLatch triggerLatch;
 
-	// ------------------------------------------------------------------------
-	//  Init & Shutdown
-	// ------------------------------------------------------------------------
-
-	@BeforeClass
-	public static void startActorSystem() {
-		actorSystem = AkkaUtils.createLocalActorSystem(new Configuration());
-	}
-
-	@AfterClass
-	public static void shutdown() {
-		actorSystem.shutdown();
-		actorSystem.awaitTermination();
-	}
-
 	@Before
 	public void createQueuesAndActors() {
 		awaitLatch = new OneShotLatch();
@@ -176,8 +153,8 @@ public class TaskAsyncCallTest {
 				mock(IOManager.class),
 				networkEnvironment,
 				mock(BroadcastVariableManager.class),
-				actorSystem.actorOf(Props.create(BlackHoleActor.class)),
-				actorSystem.actorOf(Props.create(BlackHoleActor.class)),
+				DummyActorGateway.INSTANCE,
+				DummyActorGateway.INSTANCE,
 				new FiniteDuration(60, TimeUnit.SECONDS),
 				libCache,
 				mock(FileCache.class));
@@ -236,10 +213,4 @@ public class TaskAsyncCallTest {
 			}
 		}
 	}
-	
-	public static class BlackHoleActor extends UntypedActor {
-
-		@Override
-		public void onReceive(Object message) {}
-	}
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskCancelTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskCancelTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskCancelTest.java
index 67a5f44..52154e1 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskCancelTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskCancelTest.java
@@ -25,6 +25,7 @@ import org.apache.flink.api.common.JobID;
 import org.apache.flink.configuration.ConfigConstants;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.client.JobClient;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.io.network.api.reader.RecordReader;
 import org.apache.flink.runtime.io.network.api.writer.RecordWriter;
 import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
@@ -108,7 +109,7 @@ public class TaskCancelTest {
 
 			// Run test
 			JobClient.submitJobDetached(
-					flink.jobManagerActor(), jobGraph, TestingUtils.TESTING_DURATION());
+					flink.getJobManagerGateway(), jobGraph, TestingUtils.TESTING_DURATION());
 
 			// Wait for the job to make some progress and then cancel
 			awaitRunning(
@@ -117,7 +118,7 @@ public class TaskCancelTest {
 			Thread.sleep(5000);
 
 			cancelJob(
-					flink.jobManagerActor(), jobGraph.getJobID(), TestingUtils.TESTING_DURATION());
+					flink.getJobManagerGateway(), jobGraph.getJobID(), TestingUtils.TESTING_DURATION());
 
 			// Wait for the job to be cancelled
 			JobStatus status = awaitTermination(
@@ -147,16 +148,14 @@ public class TaskCancelTest {
 	 * @param jobId The JobID of the job to cancel.
 	 * @param timeout Duration in which the JobManager must have responded.
 	 */
-	public static void cancelJob(ActorRef jobManager, JobID jobId, FiniteDuration timeout)
+	public static void cancelJob(ActorGateway jobManager, JobID jobId, FiniteDuration timeout)
 			throws Exception {
 
 		checkNotNull(jobManager);
 		checkNotNull(jobId);
 		checkNotNull(timeout);
 
-		Future<Object> ask = Patterns.ask(jobManager,
-				new CancelJob(jobId),
-				new Timeout(timeout));
+		Future<Object> ask = jobManager.ask(new CancelJob(jobId), timeout);
 
 		Object result = Await.result(ask, timeout);
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskInputSplitProviderTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskInputSplitProviderTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskInputSplitProviderTest.java
index f068978..4ccce1d 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskInputSplitProviderTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskInputSplitProviderTest.java
@@ -18,22 +18,16 @@
 
 package org.apache.flink.runtime.taskmanager;
 
-import akka.actor.ActorRef;
-import akka.actor.ActorSystem;
-import akka.actor.Props;
-import akka.actor.Status;
-import akka.actor.UntypedActor;
-import akka.testkit.JavaTestKit;
-import akka.util.Timeout;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.core.io.InputSplit;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
+import org.apache.flink.runtime.instance.BaseTestingActorGateway;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
 import org.apache.flink.runtime.messages.JobManagerMessages;
 import org.apache.flink.runtime.testingUtils.TestingUtils;
-import org.junit.AfterClass;
-import org.junit.BeforeClass;
 import org.junit.Test;
+import scala.concurrent.duration.FiniteDuration;
 
 import static org.junit.Assert.*;
 
@@ -41,31 +35,19 @@ import java.util.concurrent.TimeUnit;
 
 public class TaskInputSplitProviderTest {
 
-	private static ActorSystem system;
-
-	@BeforeClass
-	public static void setup() throws Exception {
-		system = ActorSystem.create("TestActorSystem", TestingUtils.testConfig());
-	}
-
-	@AfterClass
-	public static void teardown() throws Exception {
-		JavaTestKit.shutdownActorSystem(system);
-		system = null;
-	}
-
 	@Test
 	public void testRequestNextInputSplitWithInvalidExecutionID() {
 
 		final JobID jobID = new JobID();
 		final JobVertexID vertexID = new JobVertexID();
 		final ExecutionAttemptID executionID = new ExecutionAttemptID();
-		final Timeout timeout = new Timeout(10, TimeUnit.SECONDS);
+		final FiniteDuration timeout = new FiniteDuration(10, TimeUnit.SECONDS);
+
+		final ActorGateway gateway = new NullInputSplitGateway();
 
-		final ActorRef jobManagerRef = system.actorOf(Props.create(NullInputSplitJobManager.class));
 
 		final TaskInputSplitProvider provider = new TaskInputSplitProvider(
-				jobManagerRef,
+				gateway,
 				jobID,
 				vertexID,
 				executionID,
@@ -79,14 +61,18 @@ public class TaskInputSplitProviderTest {
 		assertTrue(nextInputSplit == null);
 	}
 
-	public static class NullInputSplitJobManager extends UntypedActor {
+	public static class NullInputSplitGateway extends BaseTestingActorGateway {
+
+		public NullInputSplitGateway() {
+			super(TestingUtils.defaultExecutionContext());
+		}
 
 		@Override
-		public void onReceive(Object message) throws Exception {
+		public Object handleMessage(Object message) throws Exception {
 			if(message instanceof JobManagerMessages.RequestNextInputSplit) {
-				sender().tell(new JobManagerMessages.NextInputSplit(null), getSelf());
+				return new JobManagerMessages.NextInputSplit(null);
 			} else {
-				sender().tell(new Status.Failure(new Exception("Invalid message type")), getSelf());
+				throw new Exception("Invalid message type");
 			}
 		}
 	}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerRegistrationTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerRegistrationTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerRegistrationTest.java
index 3e5f2cf..6d8e656 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerRegistrationTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerRegistrationTest.java
@@ -37,6 +37,7 @@ import org.apache.flink.runtime.messages.RegistrationMessages.AcknowledgeRegistr
 import org.apache.flink.runtime.messages.RegistrationMessages.RegisterTaskManager;
 import org.apache.flink.runtime.messages.RegistrationMessages.RefuseRegistration;
 import org.apache.flink.runtime.messages.TaskManagerMessages;
+import org.apache.flink.runtime.testingUtils.TestingTaskManager;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -49,6 +50,7 @@ import scala.concurrent.duration.FiniteDuration;
 import java.io.IOException;
 import java.net.InetAddress;
 import java.net.ServerSocket;
+import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
 import static org.junit.Assert.*;
@@ -148,7 +150,7 @@ public class TaskManagerRegistrationTest {
 				// start a TaskManager that tries to register at the JobManager before the JobManager is
 				// available. we give it the regular JobManager akka URL
 				final ActorRef taskManager = startTaskManager(JobManager.getLocalJobManagerAkkaURL(),
-																new Configuration());
+						new Configuration());
 				// let it try for a bit
 				Thread.sleep(6000);
 
@@ -240,7 +242,7 @@ public class TaskManagerRegistrationTest {
 						expectMsgClass(RegisterTaskManager.class);
 
 						// we decline the registration
-						getLastSender().tell(new RefuseRegistration("test reason"), getTestActor());
+						getLastSender().tell(new RefuseRegistration(UUID.randomUUID(), "test reason"), getTestActor());
 					}
 				};
 
@@ -281,17 +283,25 @@ public class TaskManagerRegistrationTest {
 				// the messages
 				final ActorRef taskManager = startTaskManager(fakeJobManager1);
 
+				final UUID leaderSessionID = UUID.randomUUID();
+
 				// validate initial registration
 				new Within(new FiniteDuration(2, TimeUnit.SECONDS)) {
 
 					@Override
 					protected void run() {
 						// the TaskManager should try to register
-						expectMsgClass(RegisterTaskManager.class);
+						RegisterTaskManager message = expectMsgClass(RegisterTaskManager.class);
 
 						// we accept the registration
-						taskManager.tell(new AcknowledgeRegistration(fakeJobManager1, new InstanceID(), 45234),
-										fakeJobManager1);
+						taskManager.tell(
+								new AcknowledgeRegistration(
+										message.registrationSessionID(),
+										leaderSessionID,
+										fakeJobManager1,
+										new InstanceID(),
+										45234),
+								fakeJobManager1);
 					}
 				};
 
@@ -331,10 +341,16 @@ public class TaskManagerRegistrationTest {
 
 					@Override
 					protected void run() {
-						expectMsgClass(RegisterTaskManager.class);
+						RegisterTaskManager message = expectMsgClass(RegisterTaskManager.class);
 
 						// we accept the registration
-						taskManager.tell(new AcknowledgeRegistration(jm2Closure, new InstanceID(), 45234),
+						taskManager.tell(
+								new AcknowledgeRegistration(
+										message.registrationSessionID(),
+										leaderSessionID,
+										jm2Closure,
+										new InstanceID(),
+										45234),
 								jm2Closure);
 					}
 				};
@@ -407,8 +423,69 @@ public class TaskManagerRegistrationTest {
 	}
 
 	@Test
-	public void testStartupWhenBlobDirectoriesAreNotWritable() {
+	public void testCheckForValidRegistrationSessionIDs() {
+		new JavaTestKit(actorSystem) {{
+			try {
+				// we make the test actor (the test kit) the JobManager to intercept
+				// the messages
+				final ActorRef taskManager = startTaskManager(getTestActor());
+
+				final UUID falseLeaderSessionID = UUID.randomUUID();
+				final UUID trueLeaderSessionID = UUID.randomUUID();
+
+				new Within(new FiniteDuration(20, TimeUnit.SECONDS)) {
+
+					@Override
+					protected void run() {
+						taskManager.tell(TaskManagerMessages.getNotifyWhenRegisteredAtJobManagerMessage(), getTestActor());
+
+						// the TaskManager should try to register
+						RegisterTaskManager registerTaskManager = expectMsgClass(RegisterTaskManager.class);
+
+						final ActorRef tm = getLastSender();
+
+						// This AcknowledgeRegistration message should be discarded because the
+						// registration session ID is wrong
+						tm.tell(new AcknowledgeRegistration(
+									UUID.randomUUID(),
+									falseLeaderSessionID,
+									getTestActor(),
+									new InstanceID(),
+									1),
+								getTestActor());
+
+						// Valid AcknowledgeRegistration message
+						tm.tell(new AcknowledgeRegistration(
+										registerTaskManager.registrationSessionID(),
+										trueLeaderSessionID,
+										getTestActor(),
+										new InstanceID(),
+										1),
+								getTestActor());
+
+						Object message = null;
+						Object confirmMessageClass = TaskManagerMessages.getRegisteredAtJobManagerMessage().getClass();
+
+						while(message == null || !(message.getClass().equals(confirmMessageClass))) {
+							message = expectMsgAnyClassOf(
+									TaskManagerMessages.getRegisteredAtJobManagerMessage().getClass(),
+									RegisterTaskManager.class,
+									TaskManagerMessages.Heartbeat.class);
+						}
+
+						tm.tell(JobManagerMessages.getRequestLeaderSessionID(), getTestActor());
+
+						expectMsgEquals(new JobManagerMessages.ResponseLeaderSessionID(Option.apply(trueLeaderSessionID)));
+					}
+				};
 
+				stopActor(taskManager);
+			}
+			catch (Throwable e) {
+				e.printStackTrace();
+				fail(e.getMessage());
+			}
+		}};
 	}
 
 	// --------------------------------------------------------------------------------------------
@@ -435,7 +512,7 @@ public class TaskManagerRegistrationTest {
 				new Some<String>(jobManagerUrl), // job manager path
 				true, // local network stack only
 				StreamingMode.BATCH_ONLY,
-				TaskManager.class);
+				TestingTaskManager.class);
 	}
 
 	private static void stopActor(ActorRef actor) {

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerTest.java
index c25b9a1..a8a852f 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerTest.java
@@ -22,7 +22,6 @@ import akka.actor.ActorRef;
 import akka.actor.ActorSystem;
 import akka.actor.Kill;
 import akka.actor.Props;
-import akka.actor.UntypedActor;
 import akka.japi.Creator;
 import akka.pattern.Patterns;
 import akka.testkit.JavaTestKit;
@@ -31,8 +30,11 @@ import akka.util.Timeout;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.configuration.ConfigConstants;
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.runtime.LeaderSessionMessageDecorator;
+import org.apache.flink.runtime.MessageDecorator;
 import org.apache.flink.runtime.StreamingMode;
 import org.apache.flink.runtime.akka.AkkaUtils;
+import org.apache.flink.runtime.akka.FlinkUntypedActor;
 import org.apache.flink.runtime.blob.BlobKey;
 import org.apache.flink.runtime.deployment.InputChannelDeploymentDescriptor;
 import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
@@ -54,6 +56,7 @@ import org.apache.flink.runtime.jobmanager.Tasks;
 import org.apache.flink.runtime.messages.Messages;
 import org.apache.flink.runtime.messages.RegistrationMessages;
 import org.apache.flink.runtime.messages.TaskManagerMessages;
+import org.apache.flink.runtime.messages.TaskMessages;
 import org.apache.flink.runtime.messages.TaskMessages.CancelTask;
 import org.apache.flink.runtime.messages.TaskMessages.PartitionState;
 import org.apache.flink.runtime.messages.TaskMessages.SubmitTask;
@@ -70,6 +73,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import scala.Option;
+import scala.Some;
 import scala.concurrent.Await;
 import scala.concurrent.Future;
 import scala.concurrent.duration.FiniteDuration;
@@ -81,9 +85,9 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
-import static org.apache.flink.runtime.messages.JobManagerMessages.ConsumerNotificationResult;
 import static org.apache.flink.runtime.messages.JobManagerMessages.RequestPartitionState;
 import static org.apache.flink.runtime.messages.JobManagerMessages.ScheduleOrUpdateConsumers;
 import static org.apache.flink.runtime.messages.TaskMessages.UpdateTaskExecutionState;
@@ -101,6 +105,10 @@ public class TaskManagerTest {
 	private static final FiniteDuration d = new FiniteDuration(20, TimeUnit.SECONDS);
 
 	private static ActorSystem system;
+
+	final static Option<UUID> leaderSessionID = new Some<UUID>(UUID.randomUUID());
+
+	final static MessageDecorator decorator = new LeaderSessionMessageDecorator(leaderSessionID);
 	
 	@BeforeClass
 	public static void setup() {
@@ -126,18 +134,25 @@ public class TaskManagerTest {
 			
 			try {
 				taskManager = createTaskManager(getTestActor(), false);
-				final ActorRef tmClosure = taskManager;
+				final ActorRef tm = taskManager;
 				
 				// handle the registration
 				new Within(d) {
 					@Override
 					protected void run() {
-						expectMsgClass(RegistrationMessages.RegisterTaskManager.class);
+						RegistrationMessages.RegisterTaskManager registrationMessage =
+								expectMsgClass(RegistrationMessages.RegisterTaskManager.class);
 						
 						final InstanceID iid = new InstanceID();
-						assertEquals(tmClosure, getLastSender());
-						tmClosure.tell(new RegistrationMessages.AcknowledgeRegistration(
-								getTestActor(), iid, 12345), getTestActor());
+						assertEquals(tm, getLastSender());
+						tm.tell(
+								new RegistrationMessages.AcknowledgeRegistration(
+										registrationMessage.registrationSessionID(),
+										leaderSessionID.get(),
+										getTestActor(),
+										iid,
+										12345),
+								getTestActor());
 					}
 				};
 
@@ -156,7 +171,7 @@ public class TaskManagerTest {
 
 					@Override
 					protected void run() {
-						tmClosure.tell(new SubmitTask(tdd), getRef());
+						tm.tell(decorator.decorate(new SubmitTask(tdd)), getRef());
 						
 						// TaskManager should acknowledge the submission
 						// heartbeats may be interleaved
@@ -169,12 +184,14 @@ public class TaskManagerTest {
 						} while (System.currentTimeMillis() < deadline);
 
 						// task should have switched to running
-						Object toRunning = new UpdateTaskExecutionState(
-								new TaskExecutionState(jid, eid, ExecutionState.RUNNING));
+						Object toRunning = decorator.decorate(
+								new TaskMessages.UpdateTaskExecutionState(
+										new TaskExecutionState(jid, eid, ExecutionState.RUNNING)));
 
 						// task should have switched to finished
-						Object toFinished = new UpdateTaskExecutionState(
-								new TaskExecutionState(jid, eid, ExecutionState.FINISHED));
+						Object toFinished = decorator.decorate(
+								new TaskMessages.UpdateTaskExecutionState(
+										new TaskExecutionState(jid, eid, ExecutionState.FINISHED)));
 						
 						deadline = System.currentTimeMillis() + 10000;
 						do {
@@ -227,7 +244,7 @@ public class TaskManagerTest {
 			ActorRef jobManager = null;
 			ActorRef taskManager = null;
 			try {
-				jobManager = system.actorOf(Props.create(SimpleJobManager.class));
+				jobManager = system.actorOf(Props.create(SimpleJobManager.class, leaderSessionID.get()));
 				taskManager = createTaskManager(jobManager, true);
 
 				final JobID jid1 = new JobID();
@@ -267,8 +284,8 @@ public class TaskManagerTest {
 									new TestingTaskManagerMessages.NotifyWhenTaskIsRunning(eid2),
 									timeout);
 
-							tm.tell(new SubmitTask(tdd1), getRef());
-							tm.tell(new SubmitTask(tdd2), getRef());
+							tm.tell(decorator.decorate(new SubmitTask(tdd1)), getRef());
+							tm.tell(decorator.decorate(new SubmitTask(tdd2)), getRef());
 
 							expectMsgEquals(Messages.getAcknowledge());
 							expectMsgEquals(Messages.getAcknowledge());
@@ -290,7 +307,7 @@ public class TaskManagerTest {
 							assertEquals(ExecutionState.RUNNING, t1.getExecutionState());
 							assertEquals(ExecutionState.RUNNING, t2.getExecutionState());
 
-							tm.tell(new CancelTask(eid1), getRef());
+							tm.tell(decorator.decorate(new CancelTask(eid1)), getRef());
 
 							expectMsgEquals(new TaskOperationResult(eid1, true));
 
@@ -306,11 +323,11 @@ public class TaskManagerTest {
 
 							assertEquals(1, runningTasks.size());
 
-							tm.tell(new CancelTask(eid1), getRef());
+							tm.tell(decorator.decorate(new CancelTask(eid1)), getRef());
 							expectMsgEquals(new TaskOperationResult(eid1, false, "No task with that execution ID was " +
 									"found."));
 
-							tm.tell(new CancelTask(eid2), getRef());
+							tm.tell(decorator.decorate(new CancelTask(eid2)), getRef());
 							expectMsgEquals(new TaskOperationResult(eid2, true));
 
 							response = Patterns.ask(tm, new TestingTaskManagerMessages.NotifyWhenTaskRemoved(eid2),
@@ -359,7 +376,7 @@ public class TaskManagerTest {
 			ActorRef jobManager = null;
 			ActorRef taskManager = null;
 			try {
-				jobManager = system.actorOf(Props.create(SimpleJobManager.class));
+				jobManager = system.actorOf(Props.create(SimpleJobManager.class, leaderSessionID.get()));
 
 				taskManager = createTaskManager(jobManager, true);
 				final ActorRef tm = taskManager;
@@ -389,8 +406,8 @@ public class TaskManagerTest {
 					@Override
 					protected void run() {
 						try {
-							tm.tell(new SubmitTask(tdd1), getRef());
-							tm.tell(new SubmitTask(tdd2), getRef());
+							tm.tell(decorator.decorate(new SubmitTask(tdd1)), getRef());
+							tm.tell(decorator.decorate(new SubmitTask(tdd2)), getRef());
 
 							expectMsgEquals(Messages.getAcknowledge());
 							expectMsgEquals(Messages.getAcknowledge());
@@ -451,7 +468,7 @@ public class TaskManagerTest {
 				final ExecutionAttemptID eid1 = new ExecutionAttemptID();
 				final ExecutionAttemptID eid2 = new ExecutionAttemptID();
 
-				jobManager = system.actorOf(Props.create(new SimpleLookupJobManagerCreator()));
+				jobManager = system.actorOf(Props.create(new SimpleLookupJobManagerCreator(leaderSessionID.get())));
 
 				taskManager = createTaskManager(jobManager, true);
 				final ActorRef tm = taskManager;
@@ -494,14 +511,14 @@ public class TaskManagerTest {
 									timeout);
 
 							// submit the sender task
-							tm.tell(new SubmitTask(tdd1), getRef());
+							tm.tell(decorator.decorate(new SubmitTask(tdd1)), getRef());
 							expectMsgEquals(Messages.getAcknowledge());
 
 							// wait until the sender task is running
 							Await.ready(t1Running, d);
 
 							// only now (after the sender is running), submit the receiver task
-							tm.tell(new SubmitTask(tdd2), getRef());
+							tm.tell(decorator.decorate(new SubmitTask(tdd2)), getRef());
 							expectMsgEquals(Messages.getAcknowledge());
 							
 							// wait until the receiver task is running
@@ -583,7 +600,9 @@ public class TaskManagerTest {
 
 				jobManager = system.actorOf(
 						Props.create(
-								new SimpleLookupFailingUpdateJobManagerCreator(eid2)
+								new SimpleLookupFailingUpdateJobManagerCreator(
+										leaderSessionID.get(),
+										eid2)
 						)
 				);
 				taskManager = createTaskManager(jobManager, true);
@@ -627,8 +646,8 @@ public class TaskManagerTest {
 									new TestingTaskManagerMessages.NotifyWhenTaskIsRunning(eid2),
 									timeout);
 
-							tm.tell(new SubmitTask(tdd2), getRef());
-							tm.tell(new SubmitTask(tdd1), getRef());
+							tm.tell(decorator.decorate(new SubmitTask(tdd2)), getRef());
+							tm.tell(decorator.decorate(new SubmitTask(tdd1)), getRef());
 
 							expectMsgEquals(Messages.getAcknowledge());
 							expectMsgEquals(Messages.getAcknowledge());
@@ -643,7 +662,7 @@ public class TaskManagerTest {
 							Task t1 = tasks.get(eid1);
 							Task t2 = tasks.get(eid2);
 
-							tm.tell(new CancelTask(eid2), getRef());
+							tm.tell(decorator.decorate(new CancelTask(eid2)), getRef());
 							expectMsgEquals(new TaskOperationResult(eid2, true));
 
 							if (t2 != null) {
@@ -654,10 +673,12 @@ public class TaskManagerTest {
 
 							if (t1 != null) {
 								if (t1.getExecutionState() == ExecutionState.RUNNING) {
-									tm.tell(new CancelTask(eid1), getRef());
+									tm.tell(decorator.decorate(new CancelTask(eid1)), getRef());
 									expectMsgEquals(new TaskOperationResult(eid1, true));
 								}
-								Future<Object> response = Patterns.ask(tm, new TestingTaskManagerMessages.NotifyWhenTaskRemoved(eid1),
+								Future<Object> response = Patterns.ask(
+										tm,
+										new TestingTaskManagerMessages.NotifyWhenTaskRemoved(eid1),
 										timeout);
 								Await.ready(response, d);
 							}
@@ -707,7 +728,7 @@ public class TaskManagerTest {
 
 				// Create the JM
 				jobManager = system.actorOf(Props.create(
-						new SimplePartitionStateLookupJobManagerCreator(resultId, getTestActor())));
+						new SimplePartitionStateLookupJobManagerCreator(leaderSessionID.get(), getTestActor())));
 
 				final int dataPort = NetUtils.getAvailablePort();
 				taskManager = createTaskManager(jobManager, true, false, dataPort);
@@ -746,7 +767,7 @@ public class TaskManagerTest {
 					@Override
 					protected void run() {
 						// Submit the task
-						tm.tell(new SubmitTask(tdd), getTestActor());
+						tm.tell(decorator.decorate(new SubmitTask(tdd)), getTestActor());
 						expectMsgClass(Messages.getAcknowledge().getClass());
 
 						// Wait to be notified about the final execution state by the mock JM
@@ -791,7 +812,7 @@ public class TaskManagerTest {
 
 				// Create the JM
 				jobManager = system.actorOf(Props.create(
-						new SimplePartitionStateLookupJobManagerCreator(resultId, getTestActor())));
+						new SimplePartitionStateLookupJobManagerCreator(leaderSessionID.get(), getTestActor())));
 
 				final int dataPort = NetUtils.getAvailablePort();
 				taskManager = createTaskManager(jobManager, true, true, dataPort);
@@ -828,7 +849,7 @@ public class TaskManagerTest {
 					@Override
 					protected void run() {
 						// Submit the task
-						tm.tell(new SubmitTask(tdd), getTestActor());
+						tm.tell(decorator.decorate(new SubmitTask(tdd)), getTestActor());
 						expectMsgClass(Messages.getAcknowledge().getClass());
 
 						// Wait to be notified about the final execution state by the mock JM
@@ -859,29 +880,56 @@ public class TaskManagerTest {
 
 	// --------------------------------------------------------------------------------------------
 
-	public static class SimpleJobManager extends UntypedActor {
+	public static class SimpleJobManager extends FlinkUntypedActor {
 
-		@Override
-		public void onReceive(Object message) throws Exception {
+		private final Option<UUID> leaderSessionID;
+
+		public SimpleJobManager(UUID leaderSessionID) {
+			this.leaderSessionID = Option.apply(leaderSessionID);
+		}
+
+		public void handleMessage(Object message) throws Exception {
 			if (message instanceof RegistrationMessages.RegisterTaskManager) {
+				final RegistrationMessages.RegisterTaskManager registerTaskManager = (RegistrationMessages.RegisterTaskManager) message;
 				final InstanceID iid = new InstanceID();
 				final ActorRef self = getSelf();
-				getSender().tell(new RegistrationMessages.AcknowledgeRegistration(self, iid, 12345), self);
+				getSender().tell(
+						decorateMessage(
+								new RegistrationMessages.AcknowledgeRegistration(
+									registerTaskManager.registrationSessionID(),
+									leaderSessionID.get(),
+									self,
+									iid,
+									12345)
+						),
+						self);
 			}
 			else if(message instanceof UpdateTaskExecutionState){
 				getSender().tell(true, getSelf());
 			}
 		}
+
+		@Override
+		protected Option<UUID> getLeaderSessionID() {
+			return leaderSessionID;
+		}
 	}
 
 	public static class SimpleLookupJobManager extends SimpleJobManager {
 
+		public SimpleLookupJobManager(UUID leaderSessionID) {
+			super(leaderSessionID);
+		}
+
 		@Override
-		public void onReceive(Object message) throws Exception {
+		public void handleMessage(Object message) throws Exception {
 			if (message instanceof ScheduleOrUpdateConsumers) {
-				getSender().tell(new ConsumerNotificationResult(true, scala.Option.<Throwable>apply(null)), getSelf());
+				getSender().tell(
+						decorateMessage(Messages.getAcknowledge()),
+						getSelf()
+						);
 			} else {
-				super.onReceive(message);
+				super.handleMessage(message);
 			}
 		}
 	}
@@ -890,12 +938,13 @@ public class TaskManagerTest {
 
 		private final Set<ExecutionAttemptID> validIDs;
 
-		public SimpleLookupFailingUpdateJobManager(Set<ExecutionAttemptID> ids) {
+		public SimpleLookupFailingUpdateJobManager(UUID leaderSessionID, Set<ExecutionAttemptID> ids) {
+			super(leaderSessionID);
 			this.validIDs = new HashSet<ExecutionAttemptID>(ids);
 		}
 
 		@Override
-		public void onReceive(Object message) throws Exception{
+		public void handleMessage(Object message) throws Exception{
 			if (message instanceof UpdateTaskExecutionState) {
 				UpdateTaskExecutionState updateMsg =
 						(UpdateTaskExecutionState) message;
@@ -906,7 +955,7 @@ public class TaskManagerTest {
 					getSender().tell(false, getSelf());
 				}
 			} else {
-				super.onReceive(message);
+				super.handleMessage(message);
 			}
 		}
 	}
@@ -915,12 +964,13 @@ public class TaskManagerTest {
 
 		private final ActorRef testActor;
 
-		public SimplePartitionStateLookupJobManager(ActorRef testActor) {
+		public SimplePartitionStateLookupJobManager(UUID leaderSessionID, ActorRef testActor) {
+			super(leaderSessionID);
 			this.testActor = testActor;
 		}
 
 		@Override
-		public void onReceive(Object message) throws Exception {
+		public void handleMessage(Object message) throws Exception {
 			if (message instanceof RequestPartitionState) {
 				final RequestPartitionState msg = (RequestPartitionState) message;
 
@@ -930,7 +980,7 @@ public class TaskManagerTest {
 						msg.partitionId().getPartitionId(),
 						ExecutionState.RUNNING);
 
-				getSender().tell(resp, getSelf());
+				getSender().tell(decorateMessage(resp), getSelf());
 			}
 			else if (message instanceof UpdateTaskExecutionState) {
 				final TaskExecutionState msg = ((UpdateTaskExecutionState) message)
@@ -940,24 +990,34 @@ public class TaskManagerTest {
 					testActor.tell(msg, self());
 				}
 			} else {
-				super.onReceive(message);
+				super.handleMessage(message);
 			}
 		}
 	}
 
 	public static class SimpleLookupJobManagerCreator implements Creator<SimpleLookupJobManager>{
 
+		private final UUID leaderSessionID;
+
+		public SimpleLookupJobManagerCreator(UUID leaderSessionID) {
+			this.leaderSessionID = leaderSessionID;
+		}
+
 		@Override
 		public SimpleLookupJobManager create() throws Exception {
-			return new SimpleLookupJobManager();
+			return new SimpleLookupJobManager(leaderSessionID);
 		}
 	}
 
 	public static class SimpleLookupFailingUpdateJobManagerCreator implements Creator<SimpleLookupFailingUpdateJobManager>{
 
+		private final UUID leaderSessionID;
+
 		private final Set<ExecutionAttemptID> validIDs;
 
-		public SimpleLookupFailingUpdateJobManagerCreator(ExecutionAttemptID ... ids) {
+		public SimpleLookupFailingUpdateJobManagerCreator(UUID leaderSessionID, ExecutionAttemptID ... ids) {
+			this.leaderSessionID = leaderSessionID;
+
 			validIDs = new HashSet<ExecutionAttemptID>();
 
 			for(ExecutionAttemptID id : ids) {
@@ -967,21 +1027,25 @@ public class TaskManagerTest {
 
 		@Override
 		public SimpleLookupFailingUpdateJobManager create() throws Exception {
-			return new SimpleLookupFailingUpdateJobManager(validIDs);
+			return new SimpleLookupFailingUpdateJobManager(leaderSessionID, validIDs);
 		}
 	}
 
 	public static class SimplePartitionStateLookupJobManagerCreator implements Creator<SimplePartitionStateLookupJobManager>{
 
+		private final UUID leaderSessionID;
+
 		private final ActorRef testActor;
 
-		public SimplePartitionStateLookupJobManagerCreator(IntermediateDataSetID dataSetId, ActorRef testActor) {
+		public SimplePartitionStateLookupJobManagerCreator(UUID leaderSessionID, ActorRef testActor) {
+			this.leaderSessionID = leaderSessionID;
+
 			this.testActor = testActor;
 		}
 
 		@Override
 		public SimplePartitionStateLookupJobManager create() throws Exception {
-			return new SimplePartitionStateLookupJobManager(testActor);
+			return new SimplePartitionStateLookupJobManager(leaderSessionID, testActor);
 		}
 	}
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java
index 4631996..c6a8cb8 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java
@@ -18,15 +18,9 @@
 
 package org.apache.flink.runtime.taskmanager;
 
-import akka.actor.ActorRef;
-import akka.actor.ActorSystem;
-import akka.actor.Kill;
-import akka.actor.Props;
-
 import com.google.common.collect.Maps;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.accumulators.AccumulatorRegistry;
-import org.apache.flink.runtime.akka.AkkaUtils;
 import org.apache.flink.runtime.blob.BlobKey;
 import org.apache.flink.runtime.broadcast.BroadcastVariableManager;
 import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
@@ -37,6 +31,7 @@ import org.apache.flink.runtime.execution.ExecutionState;
 import org.apache.flink.runtime.execution.librarycache.LibraryCacheManager;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
 import org.apache.flink.runtime.filecache.FileCache;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.io.disk.iomanager.IOManager;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.runtime.io.network.NetworkEnvironment;
@@ -51,9 +46,7 @@ import org.apache.flink.runtime.memorymanager.MemoryManager;
 import org.apache.flink.runtime.messages.TaskMessages;
 
 import org.junit.After;
-import org.junit.AfterClass;
 import org.junit.Before;
-import org.junit.BeforeClass;
 import org.junit.Test;
 import scala.concurrent.duration.FiniteDuration;
 
@@ -89,42 +82,25 @@ import static org.mockito.Mockito.when;
  */
 public class TaskTest {
 	
-	private static ActorSystem actorSystem;
-	
 	private static OneShotLatch awaitLatch;
 	private static OneShotLatch triggerLatch;
 	
-	private ActorRef taskManagerMock;
-	private ActorRef jobManagerMock;
-	private ActorRef listenerActor;
+	private ActorGateway taskManagerGateway;
+	private ActorGateway jobManagerGateway;
+	private ActorGateway listenerGateway;
 
 	private BlockingQueue<Object> taskManagerMessages;
 	private BlockingQueue<Object> jobManagerMessages;
 	private BlockingQueue<Object> listenerMessages;
 	
-	// ------------------------------------------------------------------------
-	//  Init & Shutdown
-	// ------------------------------------------------------------------------
-
-	@BeforeClass
-	public static void startActorSystem() {
-		actorSystem = AkkaUtils.createLocalActorSystem(new Configuration());
-	}
-	
-	@AfterClass
-	public static void shutdown() {
-		actorSystem.shutdown();
-		actorSystem.awaitTermination();
-	}
-	
 	@Before
 	public void createQueuesAndActors() {
 		taskManagerMessages = new LinkedBlockingQueue<Object>();
 		jobManagerMessages = new LinkedBlockingQueue<Object>();
 		listenerMessages = new LinkedBlockingQueue<Object>();
-		taskManagerMock = actorSystem.actorOf(Props.create(ForwardingActor.class, taskManagerMessages));
-		jobManagerMock = actorSystem.actorOf(Props.create(ForwardingActor.class, jobManagerMessages));
-		listenerActor = actorSystem.actorOf(Props.create(ForwardingActor.class, listenerMessages));
+		taskManagerGateway = new ForwardingActorGateway(taskManagerMessages);
+		jobManagerGateway = new ForwardingActorGateway(jobManagerMessages);
+		listenerGateway = new ForwardingActorGateway(listenerMessages);
 		
 		awaitLatch = new OneShotLatch();
 		triggerLatch = new OneShotLatch();
@@ -135,9 +111,10 @@ public class TaskTest {
 		jobManagerMessages = null;
 		taskManagerMessages = null;
 		listenerMessages = null;
-		taskManagerMock.tell(Kill.getInstance(), ActorRef.noSender());
-		jobManagerMock.tell(Kill.getInstance(), ActorRef.noSender());
-		listenerActor.tell(Kill.getInstance(), ActorRef.noSender());
+
+		taskManagerGateway = null;
+		jobManagerGateway = null;
+		listenerGateway = null;
 	}
 
 	// ------------------------------------------------------------------------
@@ -154,7 +131,7 @@ public class TaskTest {
 			assertFalse(task.isCanceledOrFailed());
 			assertNull(task.getFailureCause());
 			
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 			
 			// go into the run method. we should switch to DEPLOYING, RUNNING, then
 			// FINISHED, and all should be good
@@ -229,7 +206,7 @@ public class TaskTest {
 			assertFalse(task.isCanceledOrFailed());
 			assertNull(task.getFailureCause());
 
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			// should fail
 			task.run();
@@ -270,7 +247,7 @@ public class TaskTest {
 			
 			Task task = createTask(TestInvokableCorrect.class, libCache, network);
 
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			task.run();
 
@@ -291,7 +268,7 @@ public class TaskTest {
 	public void testInvokableInstantiationFailed() {
 		try {
 			Task task = createTask(InvokableNonInstantiable.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			task.run();
 
@@ -312,7 +289,7 @@ public class TaskTest {
 	public void testExecutionFailsInRegisterInputOutput() {
 		try {
 			Task task = createTask(InvokableWithExceptionInRegisterInOut.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			task.run();
 
@@ -333,7 +310,7 @@ public class TaskTest {
 	public void testExecutionFailsInInvoke() {
 		try {
 			Task task = createTask(InvokableWithExceptionInInvoke.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 			
 			task.run();
 
@@ -357,7 +334,7 @@ public class TaskTest {
 	public void testCancelDuringRegisterInputOutput() {
 		try {
 			Task task = createTask(InvokableBlockingInRegisterInOut.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			// run the task asynchronous
 			task.startTaskThread();
@@ -388,7 +365,7 @@ public class TaskTest {
 	public void testFailDuringRegisterInputOutput() {
 		try {
 			Task task = createTask(InvokableBlockingInRegisterInOut.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			// run the task asynchronous
 			task.startTaskThread();
@@ -419,7 +396,7 @@ public class TaskTest {
 	public void testCancelDuringInvoke() {
 		try {
 			Task task = createTask(InvokableBlockingInInvoke.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			// run the task asynchronous
 			task.startTaskThread();
@@ -453,7 +430,7 @@ public class TaskTest {
 	public void testFailExternallyDuringInvoke() {
 		try {
 			Task task = createTask(InvokableBlockingInInvoke.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			// run the task asynchronous
 			task.startTaskThread();
@@ -486,7 +463,7 @@ public class TaskTest {
 	public void testCanceledAfterExecutionFailedInRegInOut() {
 		try {
 			Task task = createTask(InvokableWithExceptionInRegisterInOut.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			task.run();
 			
@@ -510,7 +487,7 @@ public class TaskTest {
 	public void testCanceledAfterExecutionFailedInInvoke() {
 		try {
 			Task task = createTask(InvokableWithExceptionInInvoke.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			task.run();
 
@@ -537,7 +514,7 @@ public class TaskTest {
 	public void testExecutionFailesAfterCanceling() {
 		try {
 			Task task = createTask(InvokableWithExceptionOnTrigger.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			// run the task asynchronous
 			task.startTaskThread();
@@ -574,7 +551,7 @@ public class TaskTest {
 	public void testExecutionFailsAfterTaskMarkedFailed() {
 		try {
 			Task task = createTask(InvokableWithExceptionOnTrigger.class);
-			task.registerExecutionListener(listenerActor);
+			task.registerExecutionListener(listenerGateway);
 
 			// run the task asynchronous
 			task.startTaskThread();
@@ -738,15 +715,17 @@ public class TaskTest {
 		
 		TaskDeploymentDescriptor tdd = createTaskDeploymentDescriptor(invokable);
 		
-		return new Task(tdd,
-						mock(MemoryManager.class),
-						mock(IOManager.class),
-						networkEnvironment,
-						mock(BroadcastVariableManager.class),
-						taskManagerMock, jobManagerMock,
-						new FiniteDuration(60, TimeUnit.SECONDS),
-						libCache,
-						mock(FileCache.class));
+		return new Task(
+				tdd,
+				mock(MemoryManager.class),
+				mock(IOManager.class),
+				networkEnvironment,
+				mock(BroadcastVariableManager.class),
+				taskManagerGateway,
+				jobManagerGateway,
+				new FiniteDuration(60, TimeUnit.SECONDS),
+				libCache,
+				mock(FileCache.class));
 	}
 
 	private TaskDeploymentDescriptor createTaskDeploymentDescriptor(Class<? extends AbstractInvokable> invokable) {

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/akka/FlinkActorTest.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/akka/FlinkActorTest.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/akka/FlinkActorTest.scala
new file mode 100644
index 0000000..78af40e
--- /dev/null
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/akka/FlinkActorTest.scala
@@ -0,0 +1,107 @@
+/*
+ * 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.akka
+
+import java.util.UUID
+
+import akka.actor.{Props, Kill, ActorRef, ActorSystem}
+import akka.testkit.{TestActorRef, TestKit}
+import grizzled.slf4j.Logger
+import org.apache.flink.runtime.akka.FlinkUntypedActorTest.PlainRequiresLeaderSessionID
+import org.apache.flink.runtime.messages.JobManagerMessages.LeaderSessionMessage
+import org.apache.flink.runtime.{LeaderSessionMessages, FlinkActor}
+import org.apache.flink.runtime.testingUtils.TestingUtils
+import org.junit.runner.RunWith
+import org.scalatest.junit.JUnitRunner
+import org.scalatest.{FunSuiteLike, Matchers, BeforeAndAfterAll}
+
+@RunWith(classOf[JUnitRunner])
+class FlinkActorTest(_system: ActorSystem)
+  extends TestKit(_system)
+  with FunSuiteLike
+  with Matchers
+  with BeforeAndAfterAll {
+
+  def this() = this(ActorSystem("TestingActorSystem", TestingUtils.testConfig))
+
+  override def afterAll(): Unit = {
+    TestKit.shutdownActorSystem(system)
+  }
+
+  test("A Flink actor should only accept LeaderSessionMessages with a valid leader session id") {
+    val leaderSessionID = Some(UUID.randomUUID())
+    val oldSessionID = Some(UUID.randomUUID())
+
+    val props = Props(classOf[PlainFlinkActor], leaderSessionID)
+
+    val actor = TestActorRef[PlainFlinkActor](props)
+
+    actor ! LeaderSessionMessage(leaderSessionID, 1)
+    actor ! LeaderSessionMessage(oldSessionID, 1)
+    actor ! LeaderSessionMessage(leaderSessionID, 1)
+
+    actor ! 1
+
+    actor.underlyingActor.counter should be (3)
+  }
+
+  test("A Flink actor should throw an exception when receiving an unwrapped " +
+    "RequiresLeaderSessionID message") {
+    val leaderSessionID = Some(UUID.randomUUID())
+
+    val props = Props(classOf[PlainFlinkActor], leaderSessionID)
+    val actor = TestActorRef[PlainFlinkActor](props)
+
+    actor.receive(LeaderSessionMessage(leaderSessionID, 1))
+    actor.receive(1)
+
+    try {
+      actor.receive(new PlainRequiresLeaderSessionID)
+
+      fail("Expected an exception, because an RequiresLeaderSessionID message was sent to the " +
+        "FlinkActor without being wrapped in a LeaderSessionMessage.")
+    } catch {
+      case e: Exception =>
+        e.getMessage should be("Received a message PlainRequiresLeaderSessionID without a " +
+          "leader session ID, even though it requires to have one.")
+    }
+  }
+
+  def stopActor(actor: ActorRef): Unit = {
+    actor ! Kill
+  }
+
+}
+
+class PlainFlinkActor(val leaderSessionID: Option[UUID])
+  extends FlinkActor
+  with LeaderSessionMessages {
+
+  val log = Logger(getClass)
+
+  var counter = 0
+
+  /** Handle incoming messages
+    *
+    * @return
+    */
+  override def handleMessage: Receive = {
+    case _ => counter += 1
+  }
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/executiongraph/ExecutionGraphRestartTest.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/executiongraph/ExecutionGraphRestartTest.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/executiongraph/ExecutionGraphRestartTest.scala
index b3090e6..6060bc3 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/executiongraph/ExecutionGraphRestartTest.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/executiongraph/ExecutionGraphRestartTest.scala
@@ -21,7 +21,7 @@ package org.apache.flink.runtime.executiongraph
 import org.apache.flink.api.common.JobID
 import org.apache.flink.configuration.Configuration
 import org.apache.flink.runtime.akka.AkkaUtils
-import org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils.SimpleInstanceGateway
+import org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils.SimpleActorGateway
 import org.apache.flink.runtime.jobgraph.{JobStatus, JobGraph, JobVertex}
 import org.apache.flink.runtime.jobmanager.Tasks
 import org.apache.flink.runtime.jobmanager.scheduler.Scheduler
@@ -36,7 +36,7 @@ class ExecutionGraphRestartTest extends WordSpecLike with Matchers {
     "be manually restartable" in {
       try {
         val instance = ExecutionGraphTestUtils.getInstance(
-          new SimpleInstanceGateway(TestingUtils.directExecutionContext))
+          new SimpleActorGateway(TestingUtils.directExecutionContext))
 
         val scheduler = new Scheduler(TestingUtils.defaultExecutionContext)
         scheduler.newInstanceAvailable(instance)
@@ -83,7 +83,7 @@ class ExecutionGraphRestartTest extends WordSpecLike with Matchers {
     "restart itself automatically" in {
       try {
         val instance = ExecutionGraphTestUtils.getInstance(
-          new SimpleInstanceGateway(TestingUtils.directExecutionContext))
+          new SimpleActorGateway(TestingUtils.directExecutionContext))
 
         val scheduler = new Scheduler(TestingUtils.defaultExecutionContext)
         scheduler.newInstanceAvailable(instance)

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/executiongraph/TaskManagerLossFailsTasksTest.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/executiongraph/TaskManagerLossFailsTasksTest.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/executiongraph/TaskManagerLossFailsTasksTest.scala
index aaa0025..3d0d3a5 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/executiongraph/TaskManagerLossFailsTasksTest.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/executiongraph/TaskManagerLossFailsTasksTest.scala
@@ -21,7 +21,7 @@ package org.apache.flink.runtime.executiongraph
 import org.apache.flink.api.common.JobID
 import org.apache.flink.configuration.Configuration
 import org.apache.flink.runtime.akka.AkkaUtils
-import org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils.SimpleInstanceGateway
+import org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils.SimpleActorGateway
 import org.apache.flink.runtime.jobgraph.{JobStatus, JobGraph, JobVertex}
 import org.apache.flink.runtime.jobmanager.Tasks
 import org.apache.flink.runtime.jobmanager.scheduler.Scheduler
@@ -34,9 +34,9 @@ class TaskManagerLossFailsTasksTest extends WordSpecLike with Matchers {
     "fail the assigned tasks" in {
       try {
         val instance1 = ExecutionGraphTestUtils.getInstance(
-          new SimpleInstanceGateway(TestingUtils.defaultExecutionContext), 10)
+          new SimpleActorGateway(TestingUtils.defaultExecutionContext), 10)
         val instance2 = ExecutionGraphTestUtils.getInstance(
-          new SimpleInstanceGateway(TestingUtils.defaultExecutionContext), 10)
+          new SimpleActorGateway(TestingUtils.defaultExecutionContext), 10)
 
         val scheduler = new Scheduler(TestingUtils.defaultExecutionContext)
         scheduler.newInstanceAvailable(instance1)

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/CoLocationConstraintITCase.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/CoLocationConstraintITCase.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/CoLocationConstraintITCase.scala
index 19f1ea1..718496c 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/CoLocationConstraintITCase.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/CoLocationConstraintITCase.scala
@@ -25,15 +25,19 @@ import org.apache.flink.runtime.jobgraph.{JobGraph, DistributionPattern, JobVert
 import org.apache.flink.runtime.jobmanager.Tasks.{Receiver, Sender}
 import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup
 import org.apache.flink.runtime.messages.JobManagerMessages.{JobResultSuccess, SubmitJob}
-import org.apache.flink.runtime.testingUtils.TestingUtils
+import org.apache.flink.runtime.testingUtils.{ScalaTestingUtils, TestingUtils}
 import org.junit.runner.RunWith
 import org.scalatest.junit.JUnitRunner
 import org.scalatest.{BeforeAndAfterAll, Matchers, WordSpecLike}
-import scala.collection.convert.WrapAsJava
 
 @RunWith(classOf[JUnitRunner])
-class CoLocationConstraintITCase(_system: ActorSystem) extends TestKit(_system) with
-ImplicitSender with WordSpecLike with Matchers with BeforeAndAfterAll with WrapAsJava{
+class CoLocationConstraintITCase(_system: ActorSystem)
+    extends TestKit(_system)
+    with ImplicitSender
+    with WordSpecLike
+    with Matchers
+    with BeforeAndAfterAll
+    with ScalaTestingUtils {
   def this() = this(ActorSystem("TestingActorSystem", TestingUtils.testConfig))
 
   /**
@@ -64,11 +68,11 @@ ImplicitSender with WordSpecLike with Matchers with BeforeAndAfterAll with WrapA
       val jobGraph = new JobGraph("Pointwise job", sender, receiver)
 
       val cluster = TestingUtils.startTestingCluster(num_tasks)
-      val jm = cluster.getJobManager
+      val gateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          gateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
 
           expectMsgType[JobResultSuccess]


[5/6] flink git commit: [FLINK-2332] [runtime] Adds leader session IDs and registration session IDs to JobManager and TaskManager messages.

Posted by tr...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java
index f166c36..5e276bf 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java
@@ -18,8 +18,6 @@
 
 package org.apache.flink.runtime.taskmanager;
 
-import akka.actor.ActorRef;
-
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.core.fs.Path;
@@ -27,6 +25,7 @@ import org.apache.flink.runtime.accumulators.AccumulatorRegistry;
 import org.apache.flink.runtime.broadcast.BroadcastVariableManager;
 import org.apache.flink.runtime.execution.Environment;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.io.disk.iomanager.IOManager;
 import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
 import org.apache.flink.runtime.io.network.partition.consumer.InputGate;
@@ -72,25 +71,32 @@ public class RuntimeEnvironment implements Environment {
 	private final ResultPartitionWriter[] writers;
 	private final InputGate[] inputGates;
 	
-	private final ActorRef jobManagerActor;
+	private final ActorGateway jobManager;
 
 	private final AccumulatorRegistry accumulatorRegistry;
 
 	// ------------------------------------------------------------------------
 
-	public RuntimeEnvironment(JobID jobId, JobVertexID jobVertexId, ExecutionAttemptID executionId,
-								String taskName, String taskNameWithSubtasks,
-								int subtaskIndex, int parallelism,
-								Configuration jobConfiguration, Configuration taskConfiguration,
-								ClassLoader userCodeClassLoader,
-								MemoryManager memManager, IOManager ioManager,
-								BroadcastVariableManager bcVarManager,
+	public RuntimeEnvironment(
+			JobID jobId,
+			JobVertexID jobVertexId,
+			ExecutionAttemptID executionId,
+			String taskName,
+			String taskNameWithSubtasks,
+			int subtaskIndex,
+			int parallelism,
+			Configuration jobConfiguration,
+			Configuration taskConfiguration,
+			ClassLoader userCodeClassLoader,
+			MemoryManager memManager,
+			IOManager ioManager,
+			BroadcastVariableManager bcVarManager,
 								AccumulatorRegistry accumulatorRegistry,
-								InputSplitProvider splitProvider,
-								Map<String, Future<Path>> distCacheEntries,
-								ResultPartitionWriter[] writers,
-								InputGate[] inputGates,
-								ActorRef jobManagerActor) {
+			InputSplitProvider splitProvider,
+			Map<String, Future<Path>> distCacheEntries,
+			ResultPartitionWriter[] writers,
+			InputGate[] inputGates,
+			ActorGateway jobManager) {
 		
 		checkArgument(parallelism > 0 && subtaskIndex >= 0 && subtaskIndex < parallelism);
 
@@ -112,7 +118,7 @@ public class RuntimeEnvironment implements Environment {
 		this.distCacheEntries = checkNotNull(distCacheEntries);
 		this.writers = checkNotNull(writers);
 		this.inputGates = checkNotNull(inputGates);
-		this.jobManagerActor = checkNotNull(jobManagerActor);
+		this.jobManager = checkNotNull(jobManager);
 	}
 
 
@@ -238,6 +244,6 @@ public class RuntimeEnvironment implements Environment {
 		}
 		
 		AcknowledgeCheckpoint message = new AcknowledgeCheckpoint(jobId, executionId, checkpointId, serializedState);
-		jobManagerActor.tell(message, ActorRef.noSender());
+		jobManager.tell(message);
 	}
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
index 13a2ace..c4f62fb 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
@@ -18,8 +18,6 @@
 
 package org.apache.flink.runtime.taskmanager;
 
-import akka.actor.ActorRef;
-import akka.util.Timeout;
 
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.api.common.cache.DistributedCache;
@@ -37,6 +35,7 @@ import org.apache.flink.runtime.execution.ExecutionState;
 import org.apache.flink.runtime.execution.librarycache.LibraryCacheManager;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
 import org.apache.flink.runtime.filecache.FileCache;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.io.disk.iomanager.IOManager;
 import org.apache.flink.runtime.io.network.NetworkEnvironment;
 import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
@@ -163,17 +162,17 @@ public class Task implements Runnable {
 
 	private final Map<IntermediateDataSetID, SingleInputGate> inputGatesById;
 
-	/** The TaskManager actor that spawned this task */
-	private final ActorRef taskManager;
+	/** Gateway to the TaskManager that spawned this task */
+	private final ActorGateway taskManager;
 
-	/** The JobManager actor */
-	private final ActorRef jobManager;
+	/** Gateway to the JobManager */
+	private final ActorGateway jobManager;
 
 	/** All actors that want to be notified about changes in the task's execution state */
-	private final List<ActorRef> executionListenerActors;
+	private final List<ActorGateway> executionListenerActors;
 
 	/** The timeout for all ask operations on actors */
-	private final Timeout actorAskTimeout;
+	private final FiniteDuration actorAskTimeout;
 
 	/** The library cache, from which the task can request its required JAR files */
 	private final LibraryCacheManager libraryCache;
@@ -224,8 +223,8 @@ public class Task implements Runnable {
 				IOManager ioManager,
 				NetworkEnvironment networkEnvironment,
 				BroadcastVariableManager bcVarManager,
-				ActorRef taskManagerActor,
-				ActorRef jobManagerActor,
+				ActorGateway taskManagerActor,
+				ActorGateway jobManagerActor,
 				FiniteDuration actorAskTimeout,
 				LibraryCacheManager libraryCache,
 				FileCache fileCache)
@@ -254,13 +253,13 @@ public class Task implements Runnable {
 
 		this.jobManager = checkNotNull(jobManagerActor);
 		this.taskManager = checkNotNull(taskManagerActor);
-		this.actorAskTimeout = new Timeout(checkNotNull(actorAskTimeout));
+		this.actorAskTimeout = checkNotNull(actorAskTimeout);
 
 		this.libraryCache = checkNotNull(libraryCache);
 		this.fileCache = checkNotNull(fileCache);
 		this.network = checkNotNull(networkEnvironment);
 
-		this.executionListenerActors = new CopyOnWriteArrayList<ActorRef>();
+		this.executionListenerActors = new CopyOnWriteArrayList<ActorGateway>();
 
 		// create the reader and writer structures
 
@@ -568,7 +567,7 @@ public class Task implements Runnable {
 			// to know this!
 			notifyObservers(ExecutionState.RUNNING, null);
 			taskManager.tell(new TaskMessages.UpdateTaskExecutionState(
-					new TaskExecutionState(jobId, executionId, ExecutionState.RUNNING)), ActorRef.noSender());
+					new TaskExecutionState(jobId, executionId, ExecutionState.RUNNING)));
 
 			// make sure the user code classloader is accessible thread-locally
 			executingThread.setContextClassLoader(userCodeClassLoader);
@@ -750,11 +749,11 @@ public class Task implements Runnable {
 	}
 
 	private void notifyFinalState() {
-		taskManager.tell(new TaskInFinalState(executionId), ActorRef.noSender());
+		taskManager.tell(new TaskInFinalState(executionId));
 	}
 
 	private void notifyFatalError(String message, Throwable cause) {
-		taskManager.tell(new FatalError(message, cause), ActorRef.noSender());
+		taskManager.tell(new FatalError(message, cause));
 	}
 
 	// ----------------------------------------------------------------------------------------------------------------
@@ -839,14 +838,10 @@ public class Task implements Runnable {
 	//  State Listeners
 	// ------------------------------------------------------------------------
 
-	public void registerExecutionListener(ActorRef listener) {
+	public void registerExecutionListener(ActorGateway listener) {
 		executionListenerActors.add(listener);
 	}
 
-	public void unregisterExecutionListener(ActorRef listener) {
-		executionListenerActors.remove(listener);
-	}
-
 	private void notifyObservers(ExecutionState newState, Throwable error) {
 		if (error == null) {
 			LOG.info(taskNameWithSubtask + " switched to " + newState);
@@ -859,8 +854,8 @@ public class Task implements Runnable {
 		TaskMessages.UpdateTaskExecutionState actorMessage = new
 				TaskMessages.UpdateTaskExecutionState(stateUpdate);
 
-		for (ActorRef listener : executionListenerActors) {
-			listener.tell(actorMessage, ActorRef.noSender());
+		for (ActorGateway listener : executionListenerActors) {
+			listener.tell(actorMessage);
 		}
 	}
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/TaskInputSplitProvider.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/TaskInputSplitProvider.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/TaskInputSplitProvider.java
index 5a69850..cb78c16 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/TaskInputSplitProvider.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/TaskInputSplitProvider.java
@@ -18,13 +18,10 @@
 
 package org.apache.flink.runtime.taskmanager;
 
-import akka.actor.ActorRef;
-
-import akka.pattern.Patterns;
-import akka.util.Timeout;
 import org.apache.flink.core.io.InputSplit;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
 import org.apache.flink.api.common.JobID;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
 import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider;
 import org.apache.flink.runtime.messages.JobManagerMessages;
@@ -32,10 +29,11 @@ import org.apache.flink.util.InstantiationUtil;
 
 import scala.concurrent.Await;
 import scala.concurrent.Future;
+import scala.concurrent.duration.FiniteDuration;
 
 public class TaskInputSplitProvider implements InputSplitProvider {
 
-	private final ActorRef jobManager;
+	private final ActorGateway jobManager;
 	
 	private final JobID jobId;
 	
@@ -45,11 +43,15 @@ public class TaskInputSplitProvider implements InputSplitProvider {
 
 	private final ClassLoader usercodeClassLoader;
 	
-	private final Timeout timeout;
+	private final FiniteDuration timeout;
 	
-	public TaskInputSplitProvider(ActorRef jobManager, JobID jobId, JobVertexID vertexId,
-								ExecutionAttemptID executionID, ClassLoader userCodeClassLoader,
-								Timeout timeout)
+	public TaskInputSplitProvider(
+			ActorGateway jobManager,
+			JobID jobId,
+			JobVertexID vertexId,
+			ExecutionAttemptID executionID,
+			ClassLoader userCodeClassLoader,
+			FiniteDuration timeout)
 	{
 		this.jobManager = jobManager;
 		this.jobId = jobId;
@@ -62,11 +64,11 @@ public class TaskInputSplitProvider implements InputSplitProvider {
 	@Override
 	public InputSplit getNextInputSplit() {
 		try {
-			final Future<Object> response = Patterns.ask(jobManager,
+			final Future<Object> response = jobManager.ask(
 					new JobManagerMessages.RequestNextInputSplit(jobId, vertexId, executionID),
 					timeout);
 
-			final Object result = Await.result(response, timeout.duration());
+			final Object result = Await.result(response, timeout);
 
 			if(!(result instanceof JobManagerMessages.NextInputSplit)){
 				throw new RuntimeException("RequestNextInputSplit requires a response of type " +

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/ActorLogMessages.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/ActorLogMessages.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/ActorLogMessages.scala
deleted file mode 100644
index c74c339..0000000
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/ActorLogMessages.scala
+++ /dev/null
@@ -1,52 +0,0 @@
-/*
- * 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
-
-import _root_.akka.actor.Actor
-
-/**
- * Mixin to add debug message logging
- */
-trait ActorLogMessages {
-  that: Actor with ActorSynchronousLogging =>
-
-  override def receive: Receive = new Actor.Receive {
-    private val _receiveWithLogMessages = receiveWithLogMessages
-
-    override def isDefinedAt(x: Any): Boolean = _receiveWithLogMessages.isDefinedAt(x)
-
-    override def apply(x: Any): Unit = {
-      if (!log.isDebugEnabled) {
-        _receiveWithLogMessages(x)
-      }
-      else {
-        log.debug(s"Received message $x at ${that.self.path} from ${that.sender()}.")
-
-        val start = System.nanoTime()
-
-        _receiveWithLogMessages(x)
-
-        val duration = (System.nanoTime() - start) / 1000000
-        log.debug(s"Handled message $x in $duration ms from ${that.sender()}.")
-      }
-    }
-  }
-
-  def receiveWithLogMessages: Receive
-}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/ActorSynchronousLogging.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/ActorSynchronousLogging.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/ActorSynchronousLogging.scala
deleted file mode 100644
index 4d3a988..0000000
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/ActorSynchronousLogging.scala
+++ /dev/null
@@ -1,31 +0,0 @@
-/*
- * 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
-
-import _root_.akka.actor.Actor
-import grizzled.slf4j.Logger
-
-/** Adds a logger to an [[akka.actor.Actor]] implementation
-  *
-  */
-trait ActorSynchronousLogging {
-  self: Actor =>
-
-  lazy val log = Logger(getClass)
-}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/FlinkActor.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/FlinkActor.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/FlinkActor.scala
new file mode 100644
index 0000000..74a9b07
--- /dev/null
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/FlinkActor.scala
@@ -0,0 +1,48 @@
+/*
+ * 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
+
+import _root_.akka.actor.Actor
+import grizzled.slf4j.Logger
+
+/** Base trait for Flink's actors.
+  *
+  * The message handling logic is defined in the handleMessage method. This allows to mixin
+  * stackable traits which change the message receiving behaviour.
+  */
+trait FlinkActor extends Actor {
+  val log: Logger
+
+  override def receive: Receive = handleMessage
+
+  /** Handle incoming messages
+    *
+    * @return
+    */
+  def handleMessage: Receive
+
+  /** Factory method for messages. This method can be used by mixins to decorate messages
+    *
+    * @param message The message to decorate
+    * @return The decorated message
+    */
+  def decorateMessage(message: Any): Any = {
+    message
+  }
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/LeaderSessionMessageDecorator.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/LeaderSessionMessageDecorator.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/LeaderSessionMessageDecorator.scala
new file mode 100644
index 0000000..c6793ed
--- /dev/null
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/LeaderSessionMessageDecorator.scala
@@ -0,0 +1,46 @@
+/*
+ * 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
+
+import java.util.UUID
+
+import org.apache.flink.runtime.messages.JobManagerMessages.LeaderSessionMessage
+import org.apache.flink.runtime.messages.RequiresLeaderSessionID
+
+/** [[MessageDecorator]] which wraps [[RequiresLeaderSessionID]] messages in a
+  * [[LeaderSessionMessage]] with the given leader session ID.
+  *
+  * @param leaderSessionID Leader session ID which is associated with the
+  *                        [[RequiresLeaderSessionID]] message
+  */
+class LeaderSessionMessageDecorator(val leaderSessionID: Option[UUID]) extends MessageDecorator {
+
+  /** Wraps [[RequiresLeaderSessionID]] messages in a [[LeaderSessionMessage]].
+    *
+    * @param message Message to decorate
+    * @return Decorated message
+    */
+  override def decorate(message: Any): Any = {
+    message match {
+      case msg: RequiresLeaderSessionID =>
+        LeaderSessionMessage(leaderSessionID, msg)
+      case msg => msg
+    }
+  }
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/LeaderSessionMessages.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/LeaderSessionMessages.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/LeaderSessionMessages.scala
new file mode 100644
index 0000000..d54926d
--- /dev/null
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/LeaderSessionMessages.scala
@@ -0,0 +1,73 @@
+/*
+ * 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
+
+import java.util.UUID
+
+import org.apache.flink.runtime.messages.JobManagerMessages.LeaderSessionMessage
+import org.apache.flink.runtime.messages.RequiresLeaderSessionID
+
+/** Mixin to filter out [[LeaderSessionMessage]] which contain an invalid leader session id.
+  * Messages which contain a valid leader session ID are unwrapped and forwarded to the actor.
+  *
+  */
+trait LeaderSessionMessages extends FlinkActor {
+  protected def leaderSessionID: Option[UUID]
+
+  abstract override def receive: Receive = {
+    case LeaderSessionMessage(id, msg) =>
+      // Filter out messages which have not the correct leader session ID
+      (leaderSessionID, id) match {
+        case (Some(currentID), Some(msgID)) =>
+          if(currentID.equals(msgID)) {
+            // correct leader session ID
+            super.receive(msg)
+          } else {
+            // discard message because of incorrect leader session ID
+            handleDiscardedMessage(msg)
+          }
+
+        case _ => handleDiscardedMessage(msg)
+      }
+    case msg: RequiresLeaderSessionID =>
+      throw new Exception(s"Received a message $msg without a leader session ID, even though" +
+        " it requires to have one.")
+    case msg =>
+      // pass the message to the parent's receive method for further processing
+      super.receive(msg)
+  }
+
+  private def handleDiscardedMessage(msg: Any): Unit = {
+    log.debug(s"Discard message $msg because the leader session ID was not correct.")
+  }
+
+  /** Wrap [[RequiresLeaderSessionID]] messages in a [[LeaderSessionMessage]]
+    *
+    * @param message The message to decorate
+    * @return The decorated message
+    */
+  override def decorateMessage(message: Any): Any = {
+    message match {
+      case msg: RequiresLeaderSessionID =>
+        LeaderSessionMessage(leaderSessionID, super.decorateMessage(msg))
+
+      case msg => super.decorateMessage(msg)
+    }
+  }
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/LogMessages.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/LogMessages.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/LogMessages.scala
new file mode 100644
index 0000000..e1e8961
--- /dev/null
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/LogMessages.scala
@@ -0,0 +1,48 @@
+/*
+ * 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
+
+/** Mixin to add message logging if the debug log level is activated
+  *
+  */
+trait LogMessages extends FlinkActor {
+  abstract override def receive: Receive = {
+    val _receive = super.receive
+
+    new Receive {
+      override def isDefinedAt(x: Any): Boolean = _receive.isDefinedAt(x)
+
+      override def apply(x: Any): Unit = {
+        if (!log.isDebugEnabled) {
+          _receive(x)
+        }
+        else {
+          log.debug(s"Received message $x at ${context.self.path} from ${context.sender()}.")
+
+          val start = System.nanoTime()
+
+          _receive(x)
+
+          val duration = (System.nanoTime() - start) / 1000000
+          log.debug(s"Handled message $x in $duration ms from ${context.sender()}.")
+        }
+      }
+    }
+  }
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/MessageDecorator.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/MessageDecorator.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/MessageDecorator.scala
new file mode 100644
index 0000000..5b1700f
--- /dev/null
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/MessageDecorator.scala
@@ -0,0 +1,32 @@
+/*
+ * 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
+
+/** Base trait for message decorators
+  *
+  */
+trait MessageDecorator {
+
+  /** Decorates a message
+    *
+    * @param message Message to decorate
+    * @return Decorated message
+    */
+  def decorate(message: Any): Any
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/akka/AkkaUtils.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/akka/AkkaUtils.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/akka/AkkaUtils.scala
index d38e503..b8cce41 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/akka/AkkaUtils.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/akka/AkkaUtils.scala
@@ -62,8 +62,10 @@ object AkkaUtils {
    *                         parameter is None, then a local actor system will be created.
    * @return created actor system
    */
-  def createActorSystem(configuration: Configuration,
-                        listeningAddress: Option[(String, Int)]): ActorSystem = {
+  def createActorSystem(
+      configuration: Configuration,
+      listeningAddress: Option[(String, Int)])
+    : ActorSystem = {
     val akkaConfig = getAkkaConfig(configuration, listeningAddress)
     createActorSystem(akkaConfig)
   }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
index 4325e41..7bf4447 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
@@ -21,7 +21,7 @@ package org.apache.flink.runtime.jobmanager
 import java.io.{File, IOException}
 import java.lang.reflect.{InvocationTargetException, Constructor}
 import java.net.InetSocketAddress
-import java.util.Collections
+import java.util.{UUID, Collections}
 
 import akka.actor.Status.{Failure, Success}
 import akka.actor._
@@ -51,10 +51,11 @@ import org.apache.flink.runtime.taskmanager.TaskManager
 import org.apache.flink.runtime.util.ZooKeeperUtil
 import org.apache.flink.runtime.util.{SerializedValue, EnvironmentInformation}
 import org.apache.flink.runtime.webmonitor.WebMonitor
-import org.apache.flink.runtime.{StreamingMode, ActorSynchronousLogging, ActorLogMessages}
+import org.apache.flink.runtime.{FlinkActor, StreamingMode, LeaderSessionMessages}
+import org.apache.flink.runtime.{LogMessages}
 import org.apache.flink.runtime.akka.AkkaUtils
 import org.apache.flink.runtime.execution.librarycache.BlobLibraryCacheManager
-import org.apache.flink.runtime.instance.InstanceManager
+import org.apache.flink.runtime.instance.{ActorGateway, AkkaActorGateway, InstanceManager}
 import org.apache.flink.runtime.jobgraph.{JobVertexID, JobGraph, JobStatus}
 import org.apache.flink.runtime.jobmanager.scheduler.{Scheduler => FlinkScheduler}
 import org.apache.flink.runtime.messages.JobManagerMessages._
@@ -62,6 +63,7 @@ import org.apache.flink.runtime.messages.RegistrationMessages._
 import org.apache.flink.runtime.messages.TaskManagerMessages.{SendStackTrace, Heartbeat}
 import org.apache.flink.util.{ExceptionUtils, InstantiationUtil}
 
+import _root_.akka.pattern.ask
 import scala.concurrent._
 import scala.concurrent.duration._
 import scala.concurrent.forkjoin.ForkJoinPool
@@ -107,11 +109,17 @@ class JobManager(
     protected val delayBetweenRetries: Long,
     protected val timeout: FiniteDuration,
     protected val mode: StreamingMode)
-  extends Actor with ActorLogMessages with ActorSynchronousLogging {
+  extends FlinkActor
+  with LeaderSessionMessages // order of the mixin is important, we want filtering after logging
+  with LogMessages // order of the mixin is important, we want first logging
+  {
+
+  override val log = Logger(getClass)
 
   /** List of current jobs running jobs */
   protected val currentJobs = scala.collection.mutable.HashMap[JobID, (ExecutionGraph, JobInfo)]()
 
+  override val leaderSessionID = Some(UUID.randomUUID())
 
   /**
    * Run when the job manager is started. Simply logs an informational message.
@@ -125,10 +133,10 @@ class JobManager(
 
     // disconnect the registered task managers
     instanceManager.getAllRegisteredInstances.asScala.foreach {
-      _.getInstanceGateway().tell(Disconnect("JobManager is shutting down"))
+      _.getActorGateway().tell(Disconnect("JobManager is shutting down"))
     }
 
-    archive ! PoisonPill
+    archive ! decorateMessage(PoisonPill)
 
     for((e,_) <- currentJobs.values) {
       e.fail(new Exception("The JobManager is shutting down."))
@@ -151,26 +159,48 @@ class JobManager(
    *
    * @return
    */
-  override def receiveWithLogMessages: Receive = {
+  override def handleMessage: Receive = {
 
-    case RegisterTaskManager(taskManager, connectionInfo, hardwareInformation, numberOfSlots) =>
+    case RegisterTaskManager(
+      registrationSessionID,
+      taskManager,
+      connectionInfo,
+      hardwareInformation,
+      numberOfSlots) =>
 
       if (instanceManager.isRegistered(taskManager)) {
         val instanceID = instanceManager.getRegisteredInstance(taskManager).getId
 
         // IMPORTANT: Send the response to the "sender", which is not the
         //            TaskManager actor, but the ask future!
-        sender() ! AlreadyRegistered(self, instanceID, libraryCacheManager.getBlobServerPort)
+        sender() ! decorateMessage(
+          AlreadyRegistered(
+            registrationSessionID,
+            leaderSessionID.get,
+            self,
+            instanceID,
+            libraryCacheManager.getBlobServerPort)
+        )
       }
       else {
         try {
-          val instanceID = instanceManager.registerTaskManager(taskManager, connectionInfo,
-            hardwareInformation, numberOfSlots)
+          val instanceID = instanceManager.registerTaskManager(
+            taskManager,
+            connectionInfo,
+            hardwareInformation,
+            numberOfSlots,
+            leaderSessionID)
 
           // IMPORTANT: Send the response to the "sender", which is not the
           //            TaskManager actor, but the ask future!
-          sender() ! AcknowledgeRegistration(self, instanceID,
-                                             libraryCacheManager.getBlobServerPort)
+          sender() ! decorateMessage(
+            AcknowledgeRegistration(
+              registrationSessionID,
+              leaderSessionID.get,
+              self,
+              instanceID,
+              libraryCacheManager.getBlobServerPort)
+          )
 
           // to be notified when the taskManager is no longer reachable
           context.watch(taskManager)
@@ -183,15 +213,19 @@ class JobManager(
 
             // IMPORTANT: Send the response to the "sender", which is not the
             //            TaskManager actor, but the ask future!
-            sender() ! RefuseRegistration(ExceptionUtils.stringifyException(e))
+            sender() ! decorateMessage(
+              RefuseRegistration(
+                registrationSessionID,
+                ExceptionUtils.stringifyException(e))
+            )
         }
       }
 
     case RequestNumberRegisteredTaskManager =>
-      sender ! instanceManager.getNumberOfRegisteredTaskManagers
+      sender ! decorateMessage(instanceManager.getNumberOfRegisteredTaskManagers)
 
     case RequestTotalNumberOfSlots =>
-      sender ! instanceManager.getTotalNumberOfSlots
+      sender ! decorateMessage(instanceManager.getTotalNumberOfSlots)
 
     case SubmitJob(jobGraph, listen) =>
       submitJob(jobGraph, listenToEvents = listen)
@@ -206,16 +240,19 @@ class JobManager(
             executionGraph.cancel()
           }(context.dispatcher)
 
-          sender ! CancellationSuccess(jobID)
+          sender ! decorateMessage(CancellationSuccess(jobID))
         case None =>
           log.info(s"No job found with ID $jobID.")
-          sender ! CancellationFailure(jobID, new IllegalArgumentException("No job found with " +
-            s"ID $jobID."))
+          sender ! decorateMessage(
+            CancellationFailure(
+              jobID,
+              new IllegalArgumentException(s"No job found with ID $jobID."))
+          )
       }
 
     case UpdateTaskExecutionState(taskExecutionState) =>
       if (taskExecutionState == null) {
-        sender ! false
+        sender ! decorateMessage(false)
       } else {
         currentJobs.get(taskExecutionState.getJobID) match {
           case Some((executionGraph, _)) =>
@@ -223,13 +260,13 @@ class JobManager(
 
             Future {
               val result = executionGraph.updateState(taskExecutionState)
-              originalSender ! result
+              originalSender ! decorateMessage(result)
             }(context.dispatcher)
 
           case None => log.error("Cannot find execution graph for ID " +
             s"${taskExecutionState.getJobID} to change state to " +
             s"${taskExecutionState.getExecutionState}.")
-            sender ! false
+            sender ! decorateMessage(false)
         }
       }
 
@@ -283,7 +320,7 @@ class JobManager(
           null
       }
 
-      sender ! NextInputSplit(serializedInputSplit)
+      sender ! decorateMessage(NextInputSplit(serializedInputSplit))
 
     case checkpointMessage : AbstractCheckpointMessage =>
       handleCheckpointMessage(checkpointMessage)
@@ -310,20 +347,31 @@ class JobManager(
                 }
                 val result = new SerializedJobExecutionResult(jobID, jobInfo.duration,
                                                               accumulatorResults)
-                jobInfo.client ! JobResultSuccess(result)
+                jobInfo.client ! decorateMessage(JobResultSuccess(result))
 
               case JobStatus.CANCELED =>
-                jobInfo.client ! Failure(new JobCancellationException(jobID,
-                  "Job was cancelled.", error))
+                jobInfo.client ! decorateMessage(
+                  Failure(
+                    new JobCancellationException(
+                      jobID,
+                    "Job was cancelled.", error)
+                  )
+                )
 
               case JobStatus.FAILED =>
-                jobInfo.client ! Failure(new JobExecutionException(jobID,
-                  "Job execution failed.", error))
+                jobInfo.client ! decorateMessage(
+                  Failure(
+                    new JobExecutionException(
+                      jobID,
+                      "Job execution failed.",
+                      error)
+                  )
+                )
 
               case x =>
                 val exception = new JobExecutionException(jobID, s"$x is not a " +
                   "terminal state.")
-                jobInfo.client ! Failure(exception)
+                jobInfo.client ! decorateMessage(Failure(exception))
                 throw exception
             }
 
@@ -337,13 +385,17 @@ class JobManager(
     case ScheduleOrUpdateConsumers(jobId, partitionId) =>
       currentJobs.get(jobId) match {
         case Some((executionGraph, _)) =>
-          sender ! Acknowledge
+          sender ! decorateMessage(Acknowledge)
           executionGraph.scheduleOrUpdateConsumers(partitionId)
         case None =>
           log.error(s"Cannot find execution graph for job ID $jobId to schedule or update " +
             s"consumers.")
-          sender ! Failure(new IllegalStateException("Cannot find execution graph for job ID " +
-            s"$jobId to schedule or update consumers."))
+          sender ! decorateMessage(
+            Failure(
+              new IllegalStateException("Cannot find execution graph for job ID " +
+                s"$jobId to schedule or update consumers.")
+            )
+          )
       }
 
     case RequestPartitionState(jobId, partitionId, taskExecutionId, taskResultId) =>
@@ -360,15 +412,21 @@ class JobManager(
           null
       }
 
-      sender ! PartitionState(
-        taskExecutionId, taskResultId, partitionId.getPartitionId, state)
+      sender ! decorateMessage(
+        PartitionState(
+          taskExecutionId,
+          taskResultId,
+          partitionId.getPartitionId,
+          state)
+      )
 
     case RequestJobStatus(jobID) =>
       currentJobs.get(jobID) match {
-        case Some((executionGraph,_)) => sender ! CurrentJobStatus(jobID, executionGraph.getState)
+        case Some((executionGraph,_)) =>
+          sender ! decorateMessage(CurrentJobStatus(jobID, executionGraph.getState))
         case None =>
           // check the archive
-          archive forward RequestJobStatus(jobID)
+          archive forward decorateMessage(RequestJobStatus(jobID))
       }
 
     case RequestRunningJobs =>
@@ -376,16 +434,21 @@ class JobManager(
         case (_, (eg, jobInfo)) => eg
       }
 
-      sender ! RunningJobs(executionGraphs)
+      sender ! decorateMessage(RunningJobs(executionGraphs))
 
     case RequestRunningJobsStatus =>
       try {
         val jobs = currentJobs map {
-          case (_, (eg, _)) => new JobStatusMessage(eg.getJobID, eg.getJobName,
-                                            eg.getState, eg.getStatusTimestamp(JobStatus.CREATED))
+          case (_, (eg, _)) =>
+            new JobStatusMessage(
+              eg.getJobID,
+              eg.getJobName,
+              eg.getState,
+              eg.getStatusTimestamp(JobStatus.CREATED)
+            )
         }
 
-        sender ! RunningJobsStatus(jobs)
+        sender ! decorateMessage(RunningJobsStatus(jobs))
       }
       catch {
         case t: Throwable => log.error("Exception while responding to RequestRunningJobsStatus", t)
@@ -393,18 +456,22 @@ class JobManager(
 
     case RequestJob(jobID) =>
       currentJobs.get(jobID) match {
-        case Some((eg, _)) => sender ! JobFound(jobID, eg)
+        case Some((eg, _)) => sender ! decorateMessage(JobFound(jobID, eg))
         case None =>
           // check the archive
-          archive forward RequestJob(jobID)
+          archive forward decorateMessage(RequestJob(jobID))
       }
 
     case RequestBlobManagerPort =>
-      sender ! libraryCacheManager.getBlobServerPort
+      sender ! decorateMessage(libraryCacheManager.getBlobServerPort)
 
     case RequestRegisteredTaskManagers =>
       import scala.collection.JavaConverters._
-      sender ! RegisteredTaskManagers(instanceManager.getAllRegisteredInstances.asScala)
+      sender ! decorateMessage(
+        RegisteredTaskManagers(
+          instanceManager.getAllRegisteredInstances.asScala
+        )
+      )
 
     case Heartbeat(instanceID, metricsReport, accumulators) =>
       log.debug(s"Received hearbeat message from $instanceID.")
@@ -420,8 +487,8 @@ class JobManager(
     case message: InfoMessage => handleInfoRequestMessage(message, sender())
 
     case RequestStackTrace(instanceID) =>
-      val gateway = instanceManager.getRegisteredInstanceById(instanceID).getInstanceGateway
-      gateway.forward(SendStackTrace, sender())
+      val gateway = instanceManager.getRegisteredInstanceById(instanceID).getActorGateway
+      gateway.forward(SendStackTrace, new AkkaActorGateway(sender(), leaderSessionID))
 
     case Terminated(taskManager) =>
       if (instanceManager.isRegistered(taskManager)) {
@@ -432,7 +499,7 @@ class JobManager(
       }
 
     case RequestJobManagerStatus =>
-      sender() ! JobManagerStatusAlive
+      sender() ! decorateMessage(JobManagerStatusAlive)
 
     case Disconnect(msg) =>
       val taskManager = sender()
@@ -443,6 +510,9 @@ class JobManager(
         instanceManager.unregisterTaskManager(taskManager)
         context.unwatch(taskManager)
       }
+
+    case RequestLeaderSessionID =>
+      sender() ! ResponseLeaderSessionID(leaderSessionID)
   }
 
   /**
@@ -456,7 +526,11 @@ class JobManager(
    */
   private def submitJob(jobGraph: JobGraph, listenToEvents: Boolean): Unit = {
     if (jobGraph == null) {
-      sender ! Failure(new JobSubmissionException(null, "JobGraph must not be null."))
+      sender ! decorateMessage(
+        Failure(
+          new JobSubmissionException(null, "JobGraph must not be null.")
+        )
+      )
     }
     else {
       val jobId = jobGraph.getJobID
@@ -499,7 +573,8 @@ class JobManager(
             timeout,
             jobGraph.getUserJarBlobKeys,
             userCodeLoader),
-            JobInfo(sender(), System.currentTimeMillis()))
+            JobInfo(sender(), System.currentTimeMillis())
+          )
         )._1
 
         // configure the execution graph
@@ -577,23 +652,29 @@ class JobManager(
           val confirmVertices: java.util.List[ExecutionJobVertex] =
             snapshotSettings.getVerticesToConfirm.asScala.map(idToVertex).asJava
 
-          executionGraph.enableSnaphotCheckpointing(
-            snapshotSettings.getCheckpointInterval, snapshotSettings.getCheckpointTimeout,
-            triggerVertices, ackVertices, confirmVertices,
-            context.system)
+          executionGraph.enableSnapshotCheckpointing(
+            snapshotSettings.getCheckpointInterval,
+            snapshotSettings.getCheckpointTimeout,
+            triggerVertices,
+            ackVertices,
+            confirmVertices,
+            context.system,
+            leaderSessionID)
         }
 
         // get notified about job status changes
-        executionGraph.registerJobStatusListener(self)
+        executionGraph.registerJobStatusListener(new AkkaActorGateway(self, leaderSessionID))
 
         if (listenToEvents) {
           // the sender wants to be notified about state changes
-          executionGraph.registerExecutionListener(sender())
-          executionGraph.registerJobStatusListener(sender())
+          val gateway = new AkkaActorGateway(sender(), leaderSessionID)
+
+          executionGraph.registerExecutionListener(gateway)
+          executionGraph.registerJobStatusListener(gateway)
         }
 
         // done with submitting the job
-        sender ! Success(jobGraph.getJobID)
+        sender ! decorateMessage(Success(jobGraph.getJobID))
       }
       catch {
         case t: Throwable =>
@@ -612,7 +693,7 @@ class JobManager(
             new JobExecutionException(jobId, s"Failed to submit job ${jobId} (${jobName})", t)
           }
 
-          sender ! Failure(rt)
+          sender ! decorateMessage(Failure(rt))
           return
       }
 
@@ -690,14 +771,14 @@ class JobManager(
             currentJobs.get(jobID) match {
               case Some((graph, jobInfo)) =>
                 val accumulatorValues = graph.getAccumulatorsSerialized()
-                sender() ! AccumulatorResultsFound(jobID, accumulatorValues)
+          sender() ! decorateMessage(AccumulatorResultsFound(jobID, accumulatorValues))
               case None =>
                 archive.forward(message)
             }
           } catch {
           case e: Exception =>
             log.error("Cannot serialize accumulator result.", e)
-            sender() ! AccumulatorResultsErroneous(jobID, e)
+            sender() ! decorateMessage(AccumulatorResultsErroneous(jobID, e))
           }
 
         case RequestAccumulatorResultsStringified(jobId) =>
@@ -830,7 +911,7 @@ class JobManager(
         try {
           eg.prepareForArchiving()
 
-          archive ! ArchiveExecutionGraph(jobID, eg)
+          archive ! decorateMessage(ArchiveExecutionGraph(jobID, eg))
         } catch {
           case t: Throwable => log.error(s"Could not prepare the execution graph $eg for " +
             "archiving.", t)
@@ -967,11 +1048,13 @@ object JobManager {
    * @param listeningAddress The hostname where the JobManager should listen for messages.
    * @param listeningPort The port where the JobManager should listen for messages.
    */
-  def runJobManager(configuration: Configuration,
-                    executionMode: JobManagerMode,
-                    streamingMode: StreamingMode,
-                    listeningAddress: String,
-                    listeningPort: Int) : Unit = {
+  def runJobManager(
+      configuration: Configuration,
+      executionMode: JobManagerMode,
+      streamingMode: StreamingMode,
+      listeningAddress: String,
+      listeningPort: Int)
+    : Unit = {
 
     LOG.info("Starting JobManager")
 
@@ -979,8 +1062,10 @@ object JobManager {
     LOG.info(s"Starting JobManager actor system at $listeningAddress:$listeningPort.")
 
     val jobManagerSystem = try {
-      val akkaConfig = AkkaUtils.getAkkaConfig(configuration,
-                                               Some((listeningAddress, listeningPort)))
+      val akkaConfig = AkkaUtils.getAkkaConfig(
+        configuration,
+        Some((listeningAddress, listeningPort))
+      )
       if (LOG.isDebugEnabled) {
         LOG.debug("Using akka configuration\n " + akkaConfig)
       }
@@ -1003,14 +1088,20 @@ object JobManager {
     try {
       // bring up the job manager actor
       LOG.info("Starting JobManager actor")
-      val (jobManager, archiver) = startJobManagerActors(configuration, 
-                                                         jobManagerSystem, streamingMode)
+      val (jobManager, archiver) = startJobManagerActors(
+        configuration,
+        jobManagerSystem,
+        streamingMode)
 
       // start a process reaper that watches the JobManager. If the JobManager actor dies,
       // the process reaper will kill the JVM process (to ensure easy failure detection)
       LOG.debug("Starting JobManager process reaper")
       jobManagerSystem.actorOf(
-        Props(classOf[ProcessReaper], jobManager, LOG.logger, RUNTIME_FAILURE_RETURN_CODE),
+        Props(
+          classOf[ProcessReaper],
+          jobManager,
+          LOG.logger,
+          RUNTIME_FAILURE_RETURN_CODE),
         "JobManager_Process_Reaper")
 
       // bring up a local task manager, if needed
@@ -1018,16 +1109,22 @@ object JobManager {
         LOG.info("Starting embedded TaskManager for JobManager's LOCAL execution mode")
 
         val taskManagerActor = TaskManager.startTaskManagerComponentsAndActor(
-                        configuration, jobManagerSystem,
-                        listeningAddress,
-                        Some(TaskManager.TASK_MANAGER_NAME),
-                        Some(jobManager.path.toString),
-                        true, streamingMode,
-                        classOf[TaskManager])
+          configuration,
+          jobManagerSystem,
+          listeningAddress,
+          Some(TaskManager.TASK_MANAGER_NAME),
+          Some(jobManager.path.toString),
+          true,
+          streamingMode,
+          classOf[TaskManager])
 
         LOG.debug("Starting TaskManager process reaper")
         jobManagerSystem.actorOf(
-          Props(classOf[ProcessReaper], taskManagerActor, LOG.logger, RUNTIME_FAILURE_RETURN_CODE),
+          Props(
+            classOf[ProcessReaper],
+            taskManagerActor,
+            LOG.logger,
+            RUNTIME_FAILURE_RETURN_CODE),
           "TaskManager_Process_Reaper")
       }
 
@@ -1257,13 +1354,18 @@ object JobManager {
    * @param actorSystem Teh actor system running the JobManager
    * @return A tuple of references (JobManager Ref, Archiver Ref)
    */
-  def startJobManagerActors(configuration: Configuration,
-                            actorSystem: ActorSystem,
-                            streamingMode: StreamingMode): (ActorRef, ActorRef) = {
+  def startJobManagerActors(
+      configuration: Configuration,
+      actorSystem: ActorSystem,
+      streamingMode: StreamingMode)
+    : (ActorRef, ActorRef) = {
 
-    startJobManagerActors(configuration, actorSystem,
-                          Some(JOB_MANAGER_NAME), Some(ARCHIVE_NAME),
-                          streamingMode)
+    startJobManagerActors(
+      configuration,
+      actorSystem,
+      Some(JOB_MANAGER_NAME),
+      Some(ARCHIVE_NAME),
+      streamingMode)
   }
   /**
    * Starts the JobManager and job archiver based on the given configuration, in the
@@ -1279,11 +1381,13 @@ object JobManager {
    * 
    * @return A tuple of references (JobManager Ref, Archiver Ref)
    */
-  def startJobManagerActors(configuration: Configuration,
-                            actorSystem: ActorSystem,
-                            jobMangerActorName: Option[String],
-                            archiverActorName: Option[String],
-                            streamingMode: StreamingMode): (ActorRef, ActorRef) = {
+  def startJobManagerActors(
+      configuration: Configuration,
+      actorSystem: ActorSystem,
+      jobMangerActorName: Option[String],
+      archiverActorName: Option[String],
+      streamingMode: StreamingMode)
+    : (ActorRef, ActorRef) = {
 
     val (executionContext,
       instanceManager,
@@ -1352,9 +1456,11 @@ object JobManager {
     "akka://flink/user/" + JOB_MANAGER_NAME
   }
 
-  def getJobManagerRemoteReferenceFuture(address: InetSocketAddress,
-                                   system: ActorSystem,
-                                   timeout: FiniteDuration): Future[ActorRef] = {
+  def getJobManagerRemoteReferenceFuture(
+      address: InetSocketAddress,
+      system: ActorSystem,
+      timeout: FiniteDuration)
+    : Future[ActorRef] = {
 
     AkkaUtils.getReference(getRemoteJobManagerAkkaURL(address), system, timeout)
   }
@@ -1369,9 +1475,12 @@ object JobManager {
    * @return The ActorRef to the JobManager
    */
   @throws(classOf[IOException])
-  def getJobManagerRemoteReference(jobManagerUrl: String,
-                                   system: ActorSystem,
-                                   timeout: FiniteDuration): ActorRef = {
+  def getJobManagerRemoteReference(
+      jobManagerUrl: String,
+      system: ActorSystem,
+      timeout: FiniteDuration)
+    : ActorRef = {
+
     try {
       val future = AkkaUtils.getReference(jobManagerUrl, system, timeout)
       Await.result(future, timeout)
@@ -1397,9 +1506,11 @@ object JobManager {
    * @return The ActorRef to the JobManager
    */
   @throws(classOf[IOException])
-  def getJobManagerRemoteReference(address: InetSocketAddress,
-                                   system: ActorSystem,
-                                   timeout: FiniteDuration): ActorRef = {
+  def getJobManagerRemoteReference(
+      address: InetSocketAddress,
+      system: ActorSystem,
+      timeout: FiniteDuration)
+    : ActorRef = {
 
     val jmAddress = getRemoteJobManagerAkkaURL(address)
     getJobManagerRemoteReference(jmAddress, system, timeout)
@@ -1415,14 +1526,38 @@ object JobManager {
    * @return The ActorRef to the JobManager
    */
   @throws(classOf[IOException])
-  def getJobManagerRemoteReference(address: InetSocketAddress,
-                                   system: ActorSystem,
-                                   config: Configuration): ActorRef = {
+  def getJobManagerRemoteReference(
+      address: InetSocketAddress,
+      system: ActorSystem,
+      config: Configuration)
+    : ActorRef = {
 
     val timeout = AkkaUtils.getLookupTimeout(config)
     getJobManagerRemoteReference(address, system, timeout)
   }
 
+  /** Returns the [[ActorGateway]] for the provided JobManager. The function automatically
+    * retrieves the current leader session ID from the JobManager and instantiates the
+    * [[AkkaActorGateway]] with it.
+    *
+    * @param jobManager ActorRef to the [[JobManager]]
+    * @param timeout Timeout for the blocking leader session ID retrieval
+    * @throws java.lang.Exception
+    * @return Gateway to the specified JobManager
+    */
+  @throws(classOf[Exception])
+  def getJobManagerGateway(
+    jobManager: ActorRef,
+    timeout: FiniteDuration
+    ): ActorGateway = {
+    val futureLeaderSessionID = (jobManager ? RequestLeaderSessionID)(timeout)
+      .mapTo[ResponseLeaderSessionID]
+
+    val leaderSessionID = Await.result(futureLeaderSessionID, timeout).leaderSessionID
+
+    new AkkaActorGateway(jobManager, leaderSessionID)
+  }
+
   // --------------------------------------------------------------------------
   //  Utilities
   // --------------------------------------------------------------------------

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/MemoryArchivist.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/MemoryArchivist.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/MemoryArchivist.scala
index 7572e72..9f228ed 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/MemoryArchivist.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/MemoryArchivist.scala
@@ -20,13 +20,12 @@ package org.apache.flink.runtime.jobmanager
 
 import java.util
 
-import akka.actor.Actor
-
+import grizzled.slf4j.Logger
 import org.apache.flink.api.common.JobID
 import org.apache.flink.runtime.jobgraph.JobStatus
 import org.apache.flink.runtime.messages.accumulators._
+import org.apache.flink.runtime.{FlinkActor, LogMessages}
 import org.apache.flink.runtime.messages.webmonitor._
-import org.apache.flink.runtime.{ActorSynchronousLogging, ActorLogMessages}
 import org.apache.flink.runtime.executiongraph.ExecutionGraph
 import org.apache.flink.runtime.messages.ArchiveMessages._
 import org.apache.flink.runtime.messages.JobManagerMessages._
@@ -56,9 +55,11 @@ import scala.collection.mutable
  * @param max_entries Maximum number of stored Flink jobs
  */
 class MemoryArchivist(private val max_entries: Int)
-  extends Actor
-  with ActorLogMessages
-  with ActorSynchronousLogging {
+  extends FlinkActor
+  with LogMessages {
+
+  override val log = Logger(getClass)
+
   /*
    * Map of execution graphs belonging to recently started jobs with the time stamp of the last
    * received job event. The insert order is preserved through a LinkedHashMap.
@@ -70,41 +71,45 @@ class MemoryArchivist(private val max_entries: Int)
   var canceledCnt: Int = 0
   var failedCnt: Int = 0
 
-  override def receiveWithLogMessages: Receive = {
+  override def handleMessage: Receive = {
     
     /* Receive Execution Graph to archive */
     case ArchiveExecutionGraph(jobID, graph) => 
       // wrap graph inside a soft reference
       graphs.update(jobID, graph)
+
       // update job counters
       graph.getState match {
         case JobStatus.FINISHED => finishedCnt += 1
         case JobStatus.CANCELED => canceledCnt += 1
         case JobStatus.FAILED => failedCnt += 1
+          // ignore transitional states, e.g. Cancelling, Running, Failing, etc.
+        case _ =>
       }
+
       trimHistory()
 
     case RequestArchivedJob(jobID: JobID) =>
       val graph = graphs.get(jobID)
-      sender ! ArchivedJob(graph)
+      sender ! decorateMessage(ArchivedJob(graph))
 
     case RequestArchivedJobs =>
-      sender ! ArchivedJobs(graphs.values)
+      sender ! decorateMessage(ArchivedJobs(graphs.values))
 
     case RequestJob(jobID) =>
       graphs.get(jobID) match {
-        case Some(graph) => sender ! JobFound(jobID, graph)
-        case None => sender ! JobNotFound(jobID)
+        case Some(graph) => sender ! decorateMessage(JobFound(jobID, graph))
+        case None => sender ! decorateMessage(JobNotFound(jobID))
       }
 
     case RequestJobStatus(jobID) =>
       graphs.get(jobID) match {
-        case Some(graph) => sender ! CurrentJobStatus(jobID, graph.getState)
-        case None => sender ! JobNotFound(jobID)
+        case Some(graph) => sender ! decorateMessage(CurrentJobStatus(jobID, graph.getState))
+        case None => sender ! decorateMessage(JobNotFound(jobID))
       }
 
     case RequestJobCounts =>
-      sender ! (finishedCnt, canceledCnt, failedCnt)
+      sender ! decorateMessage((finishedCnt, canceledCnt, failedCnt))
 
     case _ : RequestJobsOverview =>
       try {

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/ExecutionGraphMessages.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/ExecutionGraphMessages.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/ExecutionGraphMessages.scala
index 6bc59a2..0cb3b0d 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/ExecutionGraphMessages.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/ExecutionGraphMessages.scala
@@ -45,11 +45,18 @@ object ExecutionGraphMessages {
    * @param timestamp of the execution state change
    * @param optionalMessage
    */
-  case class ExecutionStateChanged(jobID: JobID, vertexID: JobVertexID,
-                                   taskName: String, totalNumberOfSubTasks: Int, subtaskIndex: Int,
-                                   executionID: ExecutionAttemptID,
-                                   newExecutionState: ExecutionState, timestamp: Long,
-                                   optionalMessage: String){
+  case class ExecutionStateChanged(
+      jobID: JobID,
+      vertexID: JobVertexID,
+      taskName: String,
+      totalNumberOfSubTasks: Int,
+      subtaskIndex: Int,
+      executionID: ExecutionAttemptID,
+      newExecutionState: ExecutionState,
+      timestamp: Long,
+      optionalMessage: String)
+    extends RequiresLeaderSessionID {
+
     override def toString: String = {
       val oMsg = if (optionalMessage != null) {
         s"\n$optionalMessage"
@@ -69,8 +76,12 @@ object ExecutionGraphMessages {
    * @param timestamp
    * @param error
    */
-  case class JobStatusChanged(jobID: JobID, newJobStatus: JobStatus, timestamp: Long,
-                              error: Throwable){
+  case class JobStatusChanged(
+      jobID: JobID,
+      newJobStatus: JobStatus,
+      timestamp: Long,
+      error: Throwable)
+    extends RequiresLeaderSessionID {
     override def toString: String = {
       s"${timestampToString(timestamp)}\tJob execution switched to status $newJobStatus."
     }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/JobManagerMessages.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/JobManagerMessages.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/JobManagerMessages.scala
index c9a2878..e38986b 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/JobManagerMessages.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/JobManagerMessages.scala
@@ -18,6 +18,8 @@
 
 package org.apache.flink.runtime.messages
 
+import java.util.UUID
+
 import org.apache.flink.api.common.JobID
 import org.apache.flink.runtime.client.{SerializedJobExecutionResult, JobStatusMessage}
 import org.apache.flink.runtime.executiongraph.{ExecutionAttemptID, ExecutionGraph}
@@ -32,6 +34,8 @@ import scala.collection.JavaConverters._
  */
 object JobManagerMessages {
 
+  case class LeaderSessionMessage(leaderSessionID: Option[UUID], message: Any)
+
   /**
    * Submits a job to the job manager. If [[registerForEvents]] is true,
    * then the sender will be registered as listener for the state change messages.
@@ -41,6 +45,7 @@ object JobManagerMessages {
    * @param registerForEvents if true, then register for state change events
    */
   case class SubmitJob(jobGraph: JobGraph, registerForEvents: Boolean)
+    extends RequiresLeaderSessionID
 
   /**
    * Cancels a job with the given [[jobID]] at the JobManager. The result of the cancellation is
@@ -48,19 +53,22 @@ object JobManagerMessages {
    *
    * @param jobID
    */
-  case class CancelJob(jobID: JobID)
+  case class CancelJob(jobID: JobID) extends RequiresLeaderSessionID
 
   /**
    * Requesting next input split for the
    * [[org.apache.flink.runtime.executiongraph.ExecutionJobVertex]]
    * of the job specified by [[jobID]]. The next input split is sent back to the sender as a
-   * [[org.apache.flink.runtime.messages.TaskManagerMessages.NextInputSplit]] message.
+   * [[NextInputSplit]] message.
    *
    * @param jobID
    * @param vertexID
    */
-  case class RequestNextInputSplit(jobID: JobID, vertexID: JobVertexID, executionAttempt:
-  ExecutionAttemptID)
+  case class RequestNextInputSplit(
+      jobID: JobID,
+      vertexID: JobVertexID,
+      executionAttempt: ExecutionAttemptID)
+    extends RequiresLeaderSessionID
 
   /**
    * Contains the next input split for a task. This message is a response to
@@ -80,10 +88,12 @@ object JobManagerMessages {
    * @param taskExecutionId The execution attempt ID of the task requesting the partition state.
    * @param taskResultId The input gate ID of the task requesting the partition state.
    */
-  case class RequestPartitionState(jobId: JobID,
-                                   partitionId: ResultPartitionID,
-                                   taskExecutionId: ExecutionAttemptID,
-                                   taskResultId: IntermediateDataSetID)
+  case class RequestPartitionState(
+      jobId: JobID,
+      partitionId: ResultPartitionID,
+      taskExecutionId: ExecutionAttemptID,
+      taskResultId: IntermediateDataSetID)
+    extends RequiresLeaderSessionID
 
   /**
    * Notifies the [[org.apache.flink.runtime.jobmanager.JobManager]] about available data for a
@@ -101,8 +111,7 @@ object JobManagerMessages {
    * @see [[org.apache.flink.runtime.io.network.partition.ResultPartition]]
    */
   case class ScheduleOrUpdateConsumers(jobId: JobID, partitionId: ResultPartitionID)
-
-  case class ConsumerNotificationResult(success: Boolean, error: Option[Throwable] = None)
+    extends RequiresLeaderSessionID
 
   /**
    * Requests the current [[JobStatus]] of the job identified by [[jobID]]. This message triggers
@@ -142,6 +151,17 @@ object JobManagerMessages {
    */
   case object RequestBlobManagerPort
 
+  /** Requests the current leader session ID of the job manager. The result is sent back to the
+    * sender as an [[ResponseLeaderSessionID]]
+    */
+  case object RequestLeaderSessionID
+
+  /** Response to the [[RequestLeaderSessionID]] message.
+    *
+    * @param leaderSessionID
+    */
+  case class ResponseLeaderSessionID(leaderSessionID: Option[UUID])
+
   /**
    * Denotes a successful job execution.
    */
@@ -300,4 +320,8 @@ object JobManagerMessages {
   def getJobManagerStatusAlive : AnyRef = {
     JobManagerStatusAlive
   }
+
+  def getRequestLeaderSessionID: AnyRef = {
+    RequestLeaderSessionID
+  }
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/Messages.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/Messages.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/Messages.scala
index e292a47..75036b3 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/Messages.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/Messages.scala
@@ -39,7 +39,7 @@ object Messages {
    *
    * @param reason The reason for disconnecting, to be displayed in log and error messages.
    */
-  case class Disconnect(reason: String)
+  case class Disconnect(reason: String) extends RequiresLeaderSessionID
 
   /**
    * Accessor for the case object instance, to simplify Java interoperability.

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/RegistrationMessages.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/RegistrationMessages.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/RegistrationMessages.scala
index 3051d00..b435ebc 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/RegistrationMessages.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/RegistrationMessages.scala
@@ -18,6 +18,8 @@
 
 package org.apache.flink.runtime.messages
 
+import java.util.UUID
+
 import akka.actor.ActorRef
 import org.apache.flink.runtime.instance.{InstanceConnectionInfo, InstanceID, HardwareDescription}
 
@@ -32,7 +34,9 @@ object RegistrationMessages {
   /**
    * Marker trait for registration messages.
    */
-  trait RegistrationMessage
+  trait RegistrationMessage {
+    def registrationSessionID: UUID
+  }
 
   /**
    * Triggers the TaskManager to attempt a registration at the JobManager.
@@ -42,10 +46,12 @@ object RegistrationMessages {
    * @param deadline Optional deadline until when the registration must be completed.
    * @param attempt The attempt number, for logging.
    */
-  case class TriggerTaskManagerRegistration(jobManagerAkkaURL: String,
-                                            timeout: FiniteDuration,
-                                            deadline: Option[Deadline],
-                                            attempt: Int)
+  case class TriggerTaskManagerRegistration(
+      registrationSessionID: UUID,
+      jobManagerAkkaURL: String,
+      timeout: FiniteDuration,
+      deadline: Option[Deadline],
+      attempt: Int)
     extends RegistrationMessage
 
   /**
@@ -57,10 +63,12 @@ object RegistrationMessages {
    * @param resources The TaskManagers resources.
    * @param numberOfSlots The number of processing slots offered by the TaskManager.
    */
-  case class RegisterTaskManager(taskManager: ActorRef,
-                                 connectionInfo: InstanceConnectionInfo,
-                                 resources: HardwareDescription,
-                                 numberOfSlots: Int)
+  case class RegisterTaskManager(
+      registrationSessionID: UUID,
+      taskManager: ActorRef,
+      connectionInfo: InstanceConnectionInfo,
+      resources: HardwareDescription,
+      numberOfSlots: Int)
     extends RegistrationMessage
 
   /**
@@ -71,7 +79,12 @@ object RegistrationMessages {
    *                   JobManager.
    * @param blobPort The server port where the JobManager's BLOB service runs.
    */
-  case class AcknowledgeRegistration(jobManager: ActorRef, instanceID: InstanceID, blobPort: Int)
+  case class AcknowledgeRegistration(
+      registrationSessionID: UUID,
+      leaderSessionID: UUID,
+      jobManager: ActorRef,
+      instanceID: InstanceID,
+      blobPort: Int)
     extends RegistrationMessage
 
   /**
@@ -80,7 +93,12 @@ object RegistrationMessages {
    * @param instanceID The instance ID under which the TaskManager is registered.
    * @param blobPort The server port where the JobManager's BLOB service runs.
    */
-  case class AlreadyRegistered(jobManager: ActorRef, instanceID: InstanceID, blobPort: Int)
+  case class AlreadyRegistered(
+      registrationSessionID: UUID,
+      leaderSessionID: UUID,
+      jobManager: ActorRef,
+      instanceID: InstanceID,
+      blobPort: Int)
     extends RegistrationMessage
 
   /**
@@ -89,6 +107,6 @@ object RegistrationMessages {
    *
    * @param reason Reason why the task manager registration was refused
    */
-  case class RefuseRegistration(reason: String)
+  case class RefuseRegistration(registrationSessionID: UUID, reason: String)
     extends RegistrationMessage
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/TaskMessages.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/TaskMessages.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/TaskMessages.scala
index 9373576..a80ca99 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/TaskMessages.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/TaskMessages.scala
@@ -47,7 +47,7 @@ object TaskMessages {
    * @param tasks Descriptor which contains the information to start the task.
    */
   case class SubmitTask(tasks: TaskDeploymentDescriptor)
-    extends TaskMessage
+    extends TaskMessage with RequiresLeaderSessionID
 
   /**
    * Cancels the task associated with [[attemptID]]. The result is sent back to the sender as a
@@ -56,7 +56,7 @@ object TaskMessages {
    * @param attemptID The task's execution attempt ID.
    */
   case class CancelTask(attemptID: ExecutionAttemptID)
-    extends TaskMessage
+    extends TaskMessage with RequiresLeaderSessionID
 
   /**
    * Triggers a fail of specified task from the outside (as opposed to the task throwing
@@ -86,15 +86,16 @@ object TaskMessages {
    * Answer to a [[RequestPartitionState]] with the state of the respective partition.
    */
   case class PartitionState(
-    taskExecutionId: ExecutionAttemptID,
-    taskResultId: IntermediateDataSetID,
-    partitionId: IntermediateResultPartitionID,
-    state: ExecutionState) extends TaskMessage
+      taskExecutionId: ExecutionAttemptID,
+      taskResultId: IntermediateDataSetID,
+      partitionId: IntermediateResultPartitionID,
+      state: ExecutionState)
+    extends TaskMessage with RequiresLeaderSessionID
 
   /**
    * Base class for messages that update the information about location of input partitions
    */
-  abstract sealed class UpdatePartitionInfo extends TaskMessage {
+  abstract sealed class UpdatePartitionInfo extends TaskMessage with RequiresLeaderSessionID {
     def executionID: ExecutionAttemptID
   }
 
@@ -104,9 +105,10 @@ object TaskMessages {
    * @param resultId The input reader to update.
    * @param partitionInfo The partition info update.
    */
-  case class UpdateTaskSinglePartitionInfo(executionID: ExecutionAttemptID,
-                                           resultId: IntermediateDataSetID,
-                                           partitionInfo: InputChannelDeploymentDescriptor)
+  case class UpdateTaskSinglePartitionInfo(
+      executionID: ExecutionAttemptID,
+      resultId: IntermediateDataSetID,
+      partitionInfo: InputChannelDeploymentDescriptor)
     extends UpdatePartitionInfo
 
   /**
@@ -115,8 +117,8 @@ object TaskMessages {
    * @param partitionInfos List of input gates with channel descriptors to update.
    */
   case class UpdateTaskMultiplePartitionInfos(
-                    executionID: ExecutionAttemptID,
-                    partitionInfos: Seq[(IntermediateDataSetID, InputChannelDeploymentDescriptor)])
+      executionID: ExecutionAttemptID,
+      partitionInfos: Seq[(IntermediateDataSetID, InputChannelDeploymentDescriptor)])
     extends UpdatePartitionInfo
 
   /**
@@ -126,7 +128,7 @@ object TaskMessages {
    * @param executionID The task's execution attempt ID.
    */
   case class FailIntermediateResultPartitions(executionID: ExecutionAttemptID)
-    extends TaskMessage
+    extends TaskMessage with RequiresLeaderSessionID
 
 
   // --------------------------------------------------------------------------
@@ -140,7 +142,7 @@ object TaskMessages {
    * @param taskExecutionState The changed task state
    */
   case class UpdateTaskExecutionState(taskExecutionState: TaskExecutionState)
-    extends TaskMessage
+    extends TaskMessage with RequiresLeaderSessionID
 
   /**
    * Response message to updates in the task state. Send for example as a response to
@@ -152,11 +154,11 @@ object TaskMessages {
    * @param success indicating whether the operation has been successful
    * @param description Optional description for unsuccessful results.
    */
-  case class TaskOperationResult(executionID: ExecutionAttemptID,
-                                 success: Boolean,
-                                 description: String)
-    extends TaskMessage
-  {
+  case class TaskOperationResult(
+      executionID: ExecutionAttemptID,
+      success: Boolean,
+      description: String)
+    extends TaskMessage {
     def this(executionID: ExecutionAttemptID, success: Boolean) = this(executionID, success, "")
   }
 
@@ -166,10 +168,10 @@ object TaskMessages {
   // --------------------------------------------------------------------------
 
   def createUpdateTaskMultiplePartitionInfos(
-                               executionID: ExecutionAttemptID,
-                               resultIDs: java.util.List[IntermediateDataSetID],
-                               partitionInfos: java.util.List[InputChannelDeploymentDescriptor]):
-  UpdateTaskMultiplePartitionInfos = {
+      executionID: ExecutionAttemptID,
+      resultIDs: java.util.List[IntermediateDataSetID],
+      partitionInfos: java.util.List[InputChannelDeploymentDescriptor])
+    : UpdateTaskMultiplePartitionInfos = {
 
     require(resultIDs.size() == partitionInfos.size(),
       "ResultIDs must have the same length as partitionInfos.")

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/main/scala/org/apache/flink/runtime/minicluster/FlinkMiniCluster.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/minicluster/FlinkMiniCluster.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/minicluster/FlinkMiniCluster.scala
index 49c701e..6f810fc 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/minicluster/FlinkMiniCluster.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/minicluster/FlinkMiniCluster.scala
@@ -28,8 +28,11 @@ import org.apache.flink.api.common.JobSubmissionResult
 import org.apache.flink.configuration.{ConfigConstants, Configuration}
 import org.apache.flink.runtime.StreamingMode
 import org.apache.flink.runtime.akka.AkkaUtils
-import org.apache.flink.runtime.client.{JobExecutionException, JobClient, SerializedJobExecutionResult}
+import org.apache.flink.runtime.client.{JobExecutionException, JobClient,
+SerializedJobExecutionResult}
+import org.apache.flink.runtime.instance.ActorGateway
 import org.apache.flink.runtime.jobgraph.JobGraph
+import org.apache.flink.runtime.jobmanager.JobManager
 import org.apache.flink.runtime.messages.TaskManagerMessages.NotifyWhenRegisteredAtJobManager
 import org.slf4j.LoggerFactory
 
@@ -132,8 +135,9 @@ abstract class FlinkMiniCluster(
     AkkaUtils.createActorSystem(config)
   }
 
-  def getJobManager: ActorRef = {
-    jobManagerActor
+  def getJobManagerGateway(): ActorGateway = {
+    // create ActorGateway from the JobManager's ActorRef
+    JobManager.getJobManagerGateway(jobManagerActor, timeout)
   }
 
   def getTaskManagers = {
@@ -206,12 +210,17 @@ abstract class FlinkMiniCluster(
     val clientActorSystem = if (singleActorSystem) jobManagerActorSystem
     else JobClient.startJobClientActorSystem(configuration)
 
-    JobClient.submitJobAndWait(clientActorSystem, jobManagerActor, jobGraph, timeout, printUpdates)
+    JobClient.submitJobAndWait(
+      clientActorSystem,
+      getJobManagerGateway(),
+      jobGraph,
+      timeout,
+      printUpdates)
   }
 
   @throws(classOf[JobExecutionException])
   def submitJobDetached(jobGraph: JobGraph) : JobSubmissionResult = {
-    JobClient.submitJobDetached(jobManagerActor, jobGraph, timeout)
+    JobClient.submitJobDetached(getJobManagerGateway(), jobGraph, timeout)
     new JobSubmissionResult(jobGraph.getJobID)
   }
 }


[2/6] flink git commit: [FLINK-2332] [runtime] Adds leader session IDs and registration session IDs to JobManager and TaskManager messages.

Posted by tr...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/JobManagerITCase.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/JobManagerITCase.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/JobManagerITCase.scala
index 750da83..0c9f4a8 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/JobManagerITCase.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/JobManagerITCase.scala
@@ -21,14 +21,13 @@ package org.apache.flink.runtime.jobmanager
 import Tasks._
 import akka.actor.ActorSystem
 import akka.actor.Status.{Success, Failure}
-import akka.pattern.ask
 import akka.testkit.{ImplicitSender, TestKit}
 import akka.util.Timeout
 import org.apache.flink.runtime.client.JobExecutionException
 import org.apache.flink.runtime.jobgraph.{JobVertex, DistributionPattern, JobGraph, ScheduleMode}
 import org.apache.flink.runtime.messages.JobManagerMessages._
 import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages.NotifyWhenJobRemoved
-import org.apache.flink.runtime.testingUtils.TestingUtils
+import org.apache.flink.runtime.testingUtils.{ScalaTestingUtils, TestingUtils}
 import org.junit.runner.RunWith
 import org.scalatest.junit.JUnitRunner
 import org.scalatest.{BeforeAndAfterAll, Matchers, WordSpecLike}
@@ -40,8 +39,13 @@ import scala.language.postfixOps
 import scala.util.Random
 
 @RunWith(classOf[JUnitRunner])
-class JobManagerITCase(_system: ActorSystem) extends TestKit(_system) with ImplicitSender with
-WordSpecLike with Matchers with BeforeAndAfterAll {
+class JobManagerITCase(_system: ActorSystem)
+  extends TestKit(_system)
+  with ImplicitSender
+  with WordSpecLike
+  with Matchers
+  with BeforeAndAfterAll
+  with ScalaTestingUtils {
   implicit val duration = 1 minute
   implicit val timeout = Timeout.durationToTimeout(duration)
 
@@ -61,17 +65,17 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Test Job", vertex)
 
       val cluster = TestingUtils.startTestingCluster(1)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
-        val response = (jm ? RequestTotalNumberOfSlots).mapTo[Int]
+        val response = (jmGateway.ask(RequestTotalNumberOfSlots, timeout.duration)).mapTo[Int]
 
         val availableSlots = Await.result(response, duration)
 
         availableSlots should equal(1)
 
         within(2 second) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           val success = expectMsgType[Success]
 
@@ -89,7 +93,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
           }
         }
 
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       }
       finally {
@@ -106,17 +110,17 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Test Job", vertex)
 
       val cluster = TestingUtils.startTestingCluster(num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
-        val response = (jm ? RequestTotalNumberOfSlots).mapTo[Int]
+        val response = (jmGateway.ask(RequestTotalNumberOfSlots, timeout.duration)).mapTo[Int]
 
         val availableSlots = Await.result(response, duration)
 
         availableSlots should equal(num_tasks)
 
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           expectMsg(Success(jobGraph.getJobID))
           val result = expectMsgType[JobResultSuccess]
@@ -124,7 +128,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
           result.result.getJobId() should equal(jobGraph.getJobID)
         }
 
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally {
         cluster.stop()
@@ -142,11 +146,11 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       jobGraph.setAllowQueuedScheduling(true)
 
       val cluster = TestingUtils.startTestingCluster(10)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           expectMsg(Success(jobGraph.getJobID))
 
@@ -154,7 +158,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
 
           result.result.getJobId() should equal(jobGraph.getJobID)
         }
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally {
         cluster.stop()
@@ -177,11 +181,11 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Pointwise Job", sender, receiver)
 
       val cluster = TestingUtils.startTestingCluster(2 * num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           expectMsg(Success(jobGraph.getJobID))
 
@@ -189,7 +193,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
 
           result.result.getJobId() should equal(jobGraph.getJobID)
         }
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally {
         cluster.stop()
@@ -212,17 +216,17 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Bipartite Job", sender, receiver)
 
       val cluster = TestingUtils.startTestingCluster(2 * num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           expectMsg(Success(jobGraph.getJobID))
 
           expectMsgType[JobResultSuccess]
         }
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally {
         cluster.stop()
@@ -249,11 +253,11 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Bipartite Job", sender1, receiver, sender2)
 
       val cluster = TestingUtils.startTestingCluster(6 * num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           expectMsg(Success(jobGraph.getJobID))
           val failure = expectMsgType[Failure]
@@ -266,7 +270,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
           }
         }
 
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally {
         cluster.stop()
@@ -293,17 +297,17 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Bipartite Job", sender1, receiver, sender2)
 
       val cluster = TestingUtils.startTestingCluster(6 * num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
 
           expectMsgType[JobResultSuccess]
         }
 
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally {
         cluster.stop()
@@ -337,17 +341,17 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       jobGraph.setScheduleMode(ScheduleMode.ALL)
 
       val cluster = TestingUtils.startTestingCluster(num_tasks, 1)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           expectMsg(Success(jobGraph.getJobID))
 
           expectMsgType[JobResultSuccess]
 
-          jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+          jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
           expectMsg(true)
         }
       } finally {
@@ -371,16 +375,16 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Pointwise Job", sender, receiver)
 
       val cluster = TestingUtils.startTestingCluster(num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! RequestTotalNumberOfSlots
+          jmGateway.tell(RequestTotalNumberOfSlots, self)
           expectMsg(num_tasks)
         }
 
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
 
           val failure = expectMsgType[Failure]
@@ -393,7 +397,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
           }
         }
 
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally {
         cluster.stop()
@@ -419,16 +423,16 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Pointwise Job", sender, receiver)
 
       val cluster = TestingUtils.startTestingCluster(num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! RequestTotalNumberOfSlots
+          jmGateway.tell(RequestTotalNumberOfSlots, self)
           expectMsg(num_tasks)
         }
 
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
           val failure = expectMsgType[Failure]
 
@@ -440,7 +444,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
           }
         }
 
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally {
         cluster.stop()
@@ -463,11 +467,11 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Pointwise job", sender, receiver)
 
       val cluster = TestingUtils.startTestingCluster(2 * num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
           val failure = expectMsgType[Failure]
 
@@ -479,7 +483,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
           }
         }
 
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally {
         cluster.stop()
@@ -502,14 +506,14 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Pointwise job", sender, receiver)
 
       val cluster = TestingUtils.startTestingCluster(num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! RequestTotalNumberOfSlots
+          jmGateway.tell(RequestTotalNumberOfSlots, self)
           expectMsg(num_tasks)
 
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
           val failure = expectMsgType[Failure]
 
@@ -521,7 +525,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
           }
         }
 
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
 
         expectMsg(true)
       } finally {
@@ -549,14 +553,14 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Pointwise job", sender, receiver)
 
       val cluster = TestingUtils.startTestingCluster(num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! RequestTotalNumberOfSlots
+          jmGateway.tell(RequestTotalNumberOfSlots, self)
           expectMsg(num_tasks)
 
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
           val failure = expectMsgType[Failure]
 
@@ -568,7 +572,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
           }
         }
 
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally {
         cluster.stop()
@@ -591,11 +595,11 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("SubtaskInFinalStateRaceCondition", source, sink)
 
       val cluster = TestingUtils.startTestingCluster(2*num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try{
         within(TestingUtils.TESTING_DURATION){
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           expectMsg(Success(jobGraph.getJobID))
           expectMsgType[JobResultSuccess]
@@ -603,7 +607,7 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
 
         sink.finished should equal(true)
 
-        jm ! NotifyWhenJobRemoved(jobGraph.getJobID)
+        jmGateway.tell(NotifyWhenJobRemoved(jobGraph.getJobID), self)
         expectMsg(true)
       } finally{
         cluster.stop()

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/JobManagerRegistrationTest.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/JobManagerRegistrationTest.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/JobManagerRegistrationTest.scala
index 5ae6b5b..6c3580d 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/JobManagerRegistrationTest.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/JobManagerRegistrationTest.scala
@@ -19,6 +19,7 @@
 package org.apache.flink.runtime.jobmanager
 
 import java.net.InetAddress
+import java.util.UUID
 
 import akka.actor._
 import akka.testkit.{ImplicitSender, TestKit}
@@ -63,25 +64,36 @@ ImplicitSender with WordSpecLike with Matchers with BeforeAndAfterAll {
 
         var id1: InstanceID = null
         var id2: InstanceID = null
+        val registrationSessionID = UUID.randomUUID()
 
         // task manager 1
         within(1 second) {
-          jm ! RegisterTaskManager(tmDummy1, connectionInfo1, hardwareDescription, 1)
+          jm ! RegisterTaskManager(
+            registrationSessionID,
+            tmDummy1,
+            connectionInfo1,
+            hardwareDescription,
+            1)
 
           val response = receiveOne(1 second)
           response match {
-            case AcknowledgeRegistration(_, id, _) => id1 = id
+            case AcknowledgeRegistration(registrationSessionID, _,  _, id, _) => id1 = id
             case _ => fail("Wrong response message: " + response)
           }
         }
 
         // task manager 2
         within(1 second) {
-          jm ! RegisterTaskManager(tmDummy2, connectionInfo2, hardwareDescription, 1)
+          jm ! RegisterTaskManager(
+            registrationSessionID,
+            tmDummy2,
+            connectionInfo2,
+            hardwareDescription,
+            1)
 
           val response = receiveOne(1 second)
           response match {
-            case AcknowledgeRegistration(_, id, _) => id2 = id
+            case AcknowledgeRegistration(registrationSessionID, _, _, id, _) => id2 = id
             case _ => fail("Wrong response message: " + response)
           }
         }
@@ -105,11 +117,28 @@ ImplicitSender with WordSpecLike with Matchers with BeforeAndAfterAll {
       try {
         val connectionInfo = new InstanceConnectionInfo(InetAddress.getLocalHost,1)
         val hardwareDescription = HardwareDescription.extractFromSystem(10)
+
+        val registrationSessionID = UUID.randomUUID()
         
         within(1 second) {
-          jm ! RegisterTaskManager(tmDummy, connectionInfo, hardwareDescription, 1)
-          jm ! RegisterTaskManager(tmDummy, connectionInfo, hardwareDescription, 1)
-          jm ! RegisterTaskManager(tmDummy, connectionInfo, hardwareDescription, 1)
+          jm ! RegisterTaskManager(
+            registrationSessionID,
+            tmDummy,
+            connectionInfo,
+            hardwareDescription,
+            1)
+          jm ! RegisterTaskManager(
+            registrationSessionID,
+            tmDummy,
+            connectionInfo,
+            hardwareDescription,
+            1)
+          jm ! RegisterTaskManager(
+            registrationSessionID,
+            tmDummy,
+            connectionInfo,
+            hardwareDescription,
+            1)
 
           expectMsgType[AcknowledgeRegistration]
           expectMsgType[AlreadyRegistered]

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/RecoveryITCase.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/RecoveryITCase.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/RecoveryITCase.scala
index f7bf56a..0c0e28a 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/RecoveryITCase.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/RecoveryITCase.scala
@@ -19,7 +19,7 @@
 package org.apache.flink.runtime.jobmanager
 
 import akka.actor.Status.Success
-import akka.actor.{ActorRef, PoisonPill, ActorSystem}
+import akka.actor.{PoisonPill, ActorSystem}
 import akka.testkit.{ImplicitSender, TestKit}
 import org.apache.flink.configuration.{ConfigConstants, Configuration}
 import org.apache.flink.runtime.jobgraph.{JobStatus, JobGraph, DistributionPattern, JobVertex}
@@ -27,14 +27,19 @@ import org.apache.flink.runtime.jobmanager.Tasks.{BlockingOnceReceiver, FailingO
 import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup
 import org.apache.flink.runtime.messages.JobManagerMessages.{ JobResultSuccess, SubmitJob}
 import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages._
-import org.apache.flink.runtime.testingUtils.{TestingCluster, TestingUtils}
+import org.apache.flink.runtime.testingUtils.{ScalaTestingUtils, TestingCluster, TestingUtils}
 import org.junit.runner.RunWith
 import org.scalatest.{BeforeAndAfterAll, Matchers, WordSpecLike}
 import org.scalatest.junit.JUnitRunner
 
 @RunWith(classOf[JUnitRunner])
-class RecoveryITCase(_system: ActorSystem) extends TestKit(_system) with ImplicitSender with
-WordSpecLike with Matchers with BeforeAndAfterAll {
+class RecoveryITCase(_system: ActorSystem)
+  extends TestKit(_system)
+  with ImplicitSender
+  with WordSpecLike
+  with Matchers
+  with BeforeAndAfterAll
+  with ScalaTestingUtils {
 
   def this() = this(ActorSystem("TestingActorSystem", TestingUtils.testConfig))
 
@@ -74,11 +79,11 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       jobGraph.setNumberOfExecutionRetries(1)
 
       val cluster = startTestClusterWithHeartbeatTimeout(2 * NUM_TASKS, 1, "2 s")
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION){
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           expectMsg(Success(jobGraph.getJobID))
 
@@ -117,11 +122,11 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       jobGraph.setNumberOfExecutionRetries(1)
 
       val cluster = startTestClusterWithHeartbeatTimeout(NUM_TASKS, 1, "2 s")
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION){
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           expectMsg(Success(jobGraph.getJobID))
 
@@ -161,21 +166,21 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
 
       val cluster = startTestClusterWithHeartbeatTimeout(NUM_TASKS, 2, "2 s")
 
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION){
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
 
           expectMsg(Success(jobGraph.getJobID))
 
-          jm ! WaitForAllVerticesToBeRunningOrFinished(jobGraph.getJobID)
+          jmGateway.tell(WaitForAllVerticesToBeRunningOrFinished(jobGraph.getJobID), self)
 
           expectMsg(AllVerticesRunning(jobGraph.getJobID))
 
           BlockingOnceReceiver.blocking = false
-          jm ! NotifyWhenJobStatus(jobGraph.getJobID, JobStatus.RESTARTING)
-          jm ! RequestWorkingTaskManager(jobGraph.getJobID)
+          jmGateway.tell(NotifyWhenJobStatus(jobGraph.getJobID, JobStatus.RESTARTING), self)
+          jmGateway.tell(RequestWorkingTaskManager(jobGraph.getJobID), self)
 
           val WorkingTaskManager(gatewayOption) = expectMsgType[WorkingTaskManager]
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/SlotSharingITCase.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/SlotSharingITCase.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/SlotSharingITCase.scala
index 48b7971..406b054 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/SlotSharingITCase.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/SlotSharingITCase.scala
@@ -25,14 +25,19 @@ import org.apache.flink.runtime.jobgraph.{JobVertex, DistributionPattern, JobGra
 import org.apache.flink.runtime.jobmanager.Tasks.{Sender, AgnosticBinaryReceiver, Receiver}
 import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup
 import org.apache.flink.runtime.messages.JobManagerMessages.{JobResultSuccess, SubmitJob}
-import org.apache.flink.runtime.testingUtils.TestingUtils
+import org.apache.flink.runtime.testingUtils.{ScalaTestingUtils, TestingUtils}
 import org.junit.runner.RunWith
 import org.scalatest.junit.JUnitRunner
 import org.scalatest.{BeforeAndAfterAll, Matchers, WordSpecLike}
 
 @RunWith(classOf[JUnitRunner])
-class SlotSharingITCase(_system: ActorSystem) extends TestKit(_system) with ImplicitSender with
-WordSpecLike with Matchers with BeforeAndAfterAll {
+class SlotSharingITCase(_system: ActorSystem)
+  extends TestKit(_system)
+  with ImplicitSender
+  with WordSpecLike
+  with Matchers
+  with BeforeAndAfterAll
+  with ScalaTestingUtils {
   def this() = this(ActorSystem("TestingActorSystem", TestingUtils.testConfig))
 
   override def afterAll(): Unit = {
@@ -61,11 +66,11 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Pointwise Job", sender, receiver)
 
       val cluster = TestingUtils.startTestingCluster(num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
           expectMsgType[JobResultSuccess]
 
@@ -105,10 +110,11 @@ WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobGraph = new JobGraph("Bipartite job", sender1, sender2, receiver)
 
       val cluster = TestingUtils.startTestingCluster(num_tasks)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
+
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
           expectMsgType[JobResultSuccess]
         }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/TaskManagerFailsWithSlotSharingITCase.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/TaskManagerFailsWithSlotSharingITCase.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/TaskManagerFailsWithSlotSharingITCase.scala
index 0fd7e1b..7e4bf03 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/TaskManagerFailsWithSlotSharingITCase.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/jobmanager/TaskManagerFailsWithSlotSharingITCase.scala
@@ -27,14 +27,19 @@ import org.apache.flink.runtime.jobmanager.Tasks.{BlockingReceiver, Sender}
 import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup
 import org.apache.flink.runtime.messages.JobManagerMessages.SubmitJob
 import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages._
-import org.apache.flink.runtime.testingUtils.TestingUtils
+import org.apache.flink.runtime.testingUtils.{ScalaTestingUtils, TestingUtils}
 import org.junit.runner.RunWith
 import org.scalatest.junit.JUnitRunner
 import org.scalatest.{BeforeAndAfterAll, Matchers, WordSpecLike}
 
 @RunWith(classOf[JUnitRunner])
-class TaskManagerFailsWithSlotSharingITCase(_system: ActorSystem) extends TestKit(_system) with
-ImplicitSender with WordSpecLike with Matchers with BeforeAndAfterAll {
+class TaskManagerFailsWithSlotSharingITCase(_system: ActorSystem)
+  extends TestKit(_system)
+  with ImplicitSender
+  with WordSpecLike
+  with Matchers
+  with BeforeAndAfterAll
+  with ScalaTestingUtils {
 
   def this() = this(ActorSystem("TestingActorSystem", TestingUtils.testConfig))
 
@@ -64,15 +69,15 @@ ImplicitSender with WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobID = jobGraph.getJobID
 
       val cluster = TestingUtils.startTestingCluster(num_tasks/2, 2)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
       val taskManagers = cluster.getTaskManagers
 
       try{
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
 
-          jm ! WaitForAllVerticesToBeRunningOrFinished(jobID)
+          jmGateway.tell(WaitForAllVerticesToBeRunningOrFinished(jobID), self)
 
           expectMsg(AllVerticesRunning(jobID))
 
@@ -113,15 +118,15 @@ ImplicitSender with WordSpecLike with Matchers with BeforeAndAfterAll {
       val jobID = jobGraph.getJobID
 
       val cluster = TestingUtils.startTestingCluster(num_tasks/2, 2)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
       val taskManagers = cluster.getTaskManagers
 
       try{
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
 
-          jm ! WaitForAllVerticesToBeRunningOrFinished(jobID)
+          jmGateway.tell(WaitForAllVerticesToBeRunningOrFinished(jobID), self)
           expectMsg(AllVerticesRunning(jobID))
 
           //kill task manager

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/ScalaTestingUtils.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/ScalaTestingUtils.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/ScalaTestingUtils.scala
new file mode 100644
index 0000000..e63025a
--- /dev/null
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/ScalaTestingUtils.scala
@@ -0,0 +1,37 @@
+/*
+ * 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.testingUtils
+
+import akka.actor.ActorRef
+import org.apache.flink.runtime.instance.{ActorGateway, AkkaActorGateway}
+
+/** Mixing for testing utils
+  *
+  */
+trait ScalaTestingUtils {
+
+  /** Converts an [[ActorRef]] into a new [[AkkaActorGateway]] with None leader session ID
+    *
+    * @param actor ActorRef for which the ActorGateway is constructed
+    * @return [[ActorGateway]] encapsulating the given [[ActorRef]]
+    */
+  implicit def actorRef2InstanceGateway(actor: ActorRef): ActorGateway = {
+    new AkkaActorGateway(actor, None)
+  }
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingJobManager.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingJobManager.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingJobManager.scala
index ca5927d..8a7297b 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingJobManager.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingJobManager.scala
@@ -18,13 +18,13 @@
 
 package org.apache.flink.runtime.testingUtils
 
-import akka.actor.{Cancellable, Terminated, ActorRef, Props}
+import akka.actor.{Cancellable, Terminated, ActorRef}
 import akka.pattern.{ask, pipe}
 import org.apache.flink.api.common.JobID
-import org.apache.flink.runtime.ActorLogMessages
+import org.apache.flink.runtime.FlinkActor
 import org.apache.flink.runtime.execution.ExecutionState
 import org.apache.flink.runtime.jobgraph.JobStatus
-import org.apache.flink.runtime.jobmanager.{JobManager, MemoryArchivist}
+import org.apache.flink.runtime.jobmanager.JobManager
 import org.apache.flink.runtime.messages.ExecutionGraphMessages.JobStatusChanged
 import org.apache.flink.runtime.messages.Messages.Disconnect
 import org.apache.flink.runtime.messages.TaskManagerMessages.Heartbeat
@@ -32,18 +32,17 @@ import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages._
 import org.apache.flink.runtime.testingUtils.TestingMessages.DisableDisconnect
 import org.apache.flink.runtime.testingUtils.TestingTaskManagerMessages.AccumulatorsChanged
 
-import scala.collection.convert.WrapAsScala
 import scala.concurrent.Future
 import scala.concurrent.duration._
 
 import scala.language.postfixOps
 
-/**
- * Mixin for [[TestingJobManager]] to support testing messages
+/** Mixin for [[TestingJobManager]] to support testing messages
  */
-trait TestingJobManager extends ActorLogMessages with WrapAsScala {
+trait TestingJobManager extends FlinkActor {
   that: JobManager =>
 
+  import scala.collection.JavaConverters._
   import context._
 
   val waitForAllVerticesToBeRunning = scala.collection.mutable.HashMap[JobID, Set[ActorRef]]()
@@ -61,40 +60,58 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
 
   var disconnectDisabled = false
 
-  abstract override def receiveWithLogMessages: Receive = {
-    receiveTestingMessages orElse super.receiveWithLogMessages
+  abstract override def handleMessage: Receive = {
+    handleTestingMessage orElse super.handleMessage
   }
 
-  def receiveTestingMessages: Receive = {
+  def handleTestingMessage: Receive = {
     case RequestExecutionGraph(jobID) =>
       currentJobs.get(jobID) match {
-        case Some((executionGraph, jobInfo)) => sender ! ExecutionGraphFound(jobID,
-          executionGraph)
-        case None => archive.tell(RequestExecutionGraph(jobID), sender)
+        case Some((executionGraph, jobInfo)) => sender ! decorateMessage(
+          ExecutionGraphFound(
+            jobID,
+            executionGraph)
+        )
+
+        case None => archive.tell(decorateMessage(RequestExecutionGraph(jobID)), sender)
       }
 
     case WaitForAllVerticesToBeRunning(jobID) =>
       if(checkIfAllVerticesRunning(jobID)){
-        sender ! AllVerticesRunning(jobID)
+        sender ! decorateMessage(AllVerticesRunning(jobID))
       }else{
         val waiting = waitForAllVerticesToBeRunning.getOrElse(jobID, Set[ActorRef]())
         waitForAllVerticesToBeRunning += jobID -> (waiting + sender)
 
         if(periodicCheck.isEmpty){
           periodicCheck =
-            Some(context.system.scheduler.schedule(0 seconds, 200 millis, self, NotifyListeners))
+            Some(
+              context.system.scheduler.schedule(
+                0 seconds,
+                200 millis,
+                self,
+                decorateMessage(NotifyListeners)
+              )
+            )
         }
       }
     case WaitForAllVerticesToBeRunningOrFinished(jobID) =>
       if(checkIfAllVerticesRunningOrFinished(jobID)){
-        sender ! AllVerticesRunning(jobID)
+        sender ! decorateMessage(AllVerticesRunning(jobID))
       }else{
         val waiting = waitForAllVerticesToBeRunningOrFinished.getOrElse(jobID, Set[ActorRef]())
         waitForAllVerticesToBeRunningOrFinished += jobID -> (waiting + sender)
 
         if(periodicCheck.isEmpty){
           periodicCheck =
-            Some(context.system.scheduler.schedule(0 seconds, 200 millis, self, NotifyListeners))
+            Some(
+              context.system.scheduler.schedule(
+                0 seconds,
+                200 millis,
+                self,
+                decorateMessage(NotifyListeners)
+              )
+            )
         }
       }
 
@@ -110,7 +127,7 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
 
 
     case NotifyWhenJobRemoved(jobID) =>
-      val gateways = instanceManager.getAllRegisteredInstances.map(_.getInstanceGateway)
+      val gateways = instanceManager.getAllRegisteredInstances.asScala.map(_.getActorGateway)
 
       val responses = gateways.map{
         gateway => gateway.ask(NotifyWhenJobRemoved(jobID), timeout).mapTo[Boolean]
@@ -118,19 +135,19 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
 
       import context.dispatcher
 
-      Future.fold(responses)(true)(_ & _) pipeTo sender
+      Future.fold(responses)(true)(_ & _).map(decorateMessage(_)) pipeTo sender
 
     case NotifyWhenTaskManagerTerminated(taskManager) =>
       val waiting = waitForTaskManagerToBeTerminated.getOrElse(taskManager.path.name, Set())
       waitForTaskManagerToBeTerminated += taskManager.path.name -> (waiting + sender)
 
     case msg@Terminated(taskManager) =>
-      super.receiveWithLogMessages(msg)
+      super.handleMessage(msg)
 
       waitForTaskManagerToBeTerminated.remove(taskManager.path.name) foreach {
         _ foreach {
           listener =>
-            listener ! TaskManagerTerminated(taskManager)
+            listener ! decorateMessage(TaskManagerTerminated(taskManager))
         }
       }
 
@@ -157,7 +174,7 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
      * we previously received an [[AccumulatorsChanged]] message.
      */
     case msg : Heartbeat =>
-      super.receiveWithLogMessages(msg)
+      super.handleMessage(msg)
 
       waitForAccumulatorUpdate foreach {
         case (jobID, (updated, actors)) if updated =>
@@ -177,18 +194,22 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
     case RequestWorkingTaskManager(jobID) =>
       currentJobs.get(jobID) match {
         case Some((eg, _)) =>
-          if(eg.getAllExecutionVertices.isEmpty){
-            sender ! WorkingTaskManager(None)
+          if(eg.getAllExecutionVertices.asScala.isEmpty){
+            sender ! decorateMessage(WorkingTaskManager(None))
           } else {
-            val resource = eg.getAllExecutionVertices.head.getCurrentAssignedResource
+            val resource = eg.getAllExecutionVertices.asScala.head.getCurrentAssignedResource
 
             if(resource == null){
-              sender ! WorkingTaskManager(None)
+              sender ! decorateMessage(WorkingTaskManager(None))
             } else {
-              sender ! WorkingTaskManager(Some(resource.getInstance().getInstanceGateway))
+              sender ! decorateMessage(
+                WorkingTaskManager(
+                  Some(resource.getInstance().getActorGateway)
+                )
+              )
             }
           }
-        case None => sender ! WorkingTaskManager(None)
+        case None => sender ! decorateMessage(WorkingTaskManager(None))
       }
 
 
@@ -201,14 +222,14 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
       jobStatusListener += state -> (listener + sender)
 
     case msg@JobStatusChanged(jobID, newJobStatus, _, _) =>
-      super.receiveWithLogMessages(msg)
+      super.handleMessage(msg)
 
       val cleanup = waitForJobStatus.get(jobID) match {
         case Some(stateListener) =>
           stateListener.remove(newJobStatus) match {
             case Some(listeners) =>
               listeners foreach {
-                _ ! JobStatusIs(jobID, newJobStatus)
+                _ ! decorateMessage(JobStatusIs(jobID, newJobStatus))
               }
             case _ =>
           }
@@ -226,14 +247,14 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
 
     case msg: Disconnect =>
       if (!disconnectDisabled) {
-        super.receiveWithLogMessages(msg)
+        super.handleMessage(msg)
 
         val taskManager = sender
 
         waitForTaskManagerToBeTerminated.remove(taskManager.path.name) foreach {
           _ foreach {
             listener =>
-              listener ! TaskManagerTerminated(taskManager)
+              listener ! decorateMessage(TaskManagerTerminated(taskManager))
           }
         }
       }
@@ -242,7 +263,7 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
   def checkIfAllVerticesRunning(jobID: JobID): Boolean = {
     currentJobs.get(jobID) match {
       case Some((eg, _)) =>
-        eg.getAllExecutionVertices.forall( _.getExecutionState == ExecutionState.RUNNING)
+        eg.getAllExecutionVertices.asScala.forall( _.getExecutionState == ExecutionState.RUNNING)
       case None => false
     }
   }
@@ -250,7 +271,7 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
   def checkIfAllVerticesRunningOrFinished(jobID: JobID): Boolean = {
     currentJobs.get(jobID) match {
       case Some((eg, _)) =>
-        eg.getAllExecutionVertices.forall {
+        eg.getAllExecutionVertices.asScala.forall {
           case vertex =>
             (vertex.getExecutionState == ExecutionState.RUNNING
               || vertex.getExecutionState == ExecutionState.FINISHED)
@@ -264,7 +285,7 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
       waitForAllVerticesToBeRunning.remove(jobID) match {
         case Some(listeners) =>
           for (listener <- listeners) {
-            listener ! AllVerticesRunning(jobID)
+            listener ! decorateMessage(AllVerticesRunning(jobID))
           }
         case _ =>
       }
@@ -274,7 +295,7 @@ trait TestingJobManager extends ActorLogMessages with WrapAsScala {
       waitForAllVerticesToBeRunningOrFinished.remove(jobID) match {
         case Some(listeners) =>
           for (listener <- listeners) {
-            listener ! AllVerticesRunning(jobID)
+            listener ! decorateMessage(AllVerticesRunning(jobID))
           }
         case _ =>
       }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingJobManagerMessages.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingJobManagerMessages.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingJobManagerMessages.scala
index bbf7b2d..17beff0 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingJobManagerMessages.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingJobManagerMessages.scala
@@ -23,7 +23,7 @@ import org.apache.flink.api.common.JobID
 import org.apache.flink.api.common.accumulators.Accumulator
 import org.apache.flink.runtime.accumulators.AccumulatorRegistry
 import org.apache.flink.runtime.executiongraph.{ExecutionAttemptID, ExecutionGraph}
-import org.apache.flink.runtime.instance.InstanceGateway
+import org.apache.flink.runtime.instance.ActorGateway
 import org.apache.flink.runtime.jobgraph.JobStatus
 import java.util.Map
 
@@ -47,7 +47,7 @@ object TestingJobManagerMessages {
   case class NotifyWhenJobRemoved(jobID: JobID)
 
   case class RequestWorkingTaskManager(jobID: JobID)
-  case class WorkingTaskManager(gatewayOption: Option[InstanceGateway])
+  case class WorkingTaskManager(gatewayOption: Option[ActorGateway])
 
   case class NotifyWhenJobStatus(jobID: JobID, state: JobStatus)
   case class JobStatusIs(jobID: JobID, state: JobStatus)

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingMemoryArchivist.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingMemoryArchivist.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingMemoryArchivist.scala
index b8d1217..f105ccc 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingMemoryArchivist.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingMemoryArchivist.scala
@@ -18,27 +18,27 @@
 
 package org.apache.flink.runtime.testingUtils
 
-import org.apache.flink.runtime.ActorLogMessages
+import org.apache.flink.runtime.{FlinkActor}
 import org.apache.flink.runtime.jobmanager.MemoryArchivist
 import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages.{ExecutionGraphNotFound, ExecutionGraphFound, RequestExecutionGraph}
 
 /**
  * Mixin for the [[MemoryArchivist]] to support testing messages
  */
-trait TestingMemoryArchivist extends ActorLogMessages {
+trait TestingMemoryArchivist extends FlinkActor {
   self: MemoryArchivist =>
 
-  abstract override def receiveWithLogMessages = {
-    receiveTestingMessages orElse super.receiveWithLogMessages
+  abstract override def handleMessage: Receive = {
+    handleTestingMessage orElse super.handleMessage
   }
 
-  def receiveTestingMessages: Receive = {
+  def handleTestingMessage: Receive = {
     case RequestExecutionGraph(jobID) =>
       val executionGraph = graphs.get(jobID)
       
       executionGraph match {
-        case Some(graph) => sender ! ExecutionGraphFound(jobID, graph)
-        case None => sender ! ExecutionGraphNotFound(jobID)
+        case Some(graph) => sender ! decorateMessage(ExecutionGraphFound(jobID, graph))
+        case None => sender ! decorateMessage(ExecutionGraphNotFound(jobID))
       }
   }
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingTaskManager.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingTaskManager.scala b/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingTaskManager.scala
index 6f8ddd3..8cc1e92 100644
--- a/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingTaskManager.scala
+++ b/flink-runtime/src/test/scala/org/apache/flink/runtime/testingUtils/TestingTaskManager.scala
@@ -26,8 +26,10 @@ import org.apache.flink.runtime.instance.InstanceConnectionInfo
 import org.apache.flink.runtime.io.disk.iomanager.IOManager
 import org.apache.flink.runtime.io.network.NetworkEnvironment
 import org.apache.flink.runtime.memorymanager.DefaultMemoryManager
+import org.apache.flink.runtime.messages.JobManagerMessages.{ResponseLeaderSessionID,
+RequestLeaderSessionID}
 import org.apache.flink.runtime.messages.Messages.Disconnect
-import org.apache.flink.runtime.messages.TaskMessages.{TaskInFinalState, UpdateTaskExecutionState}
+import org.apache.flink.runtime.messages.TaskMessages.{UpdateTaskExecutionState, TaskInFinalState}
 import org.apache.flink.runtime.taskmanager.{TaskManagerConfiguration, TaskManager}
 import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages.NotifyWhenJobRemoved
 import org.apache.flink.runtime.testingUtils.TestingMessages.DisableDisconnect
@@ -36,22 +38,35 @@ import org.apache.flink.runtime.testingUtils.TestingTaskManagerMessages._
 import scala.concurrent.duration._
 import scala.language.postfixOps
 
-/**
- * Subclass of the [[TaskManager]] to support testing messages
- */
-class TestingTaskManager(config: TaskManagerConfiguration,
-                         connectionInfo: InstanceConnectionInfo,
-                         jobManagerAkkaURL: String,
-                         memoryManager: DefaultMemoryManager,
-                         ioManager: IOManager,
-                         network: NetworkEnvironment,
-                         numberOfSlots: Int)
-  extends TaskManager(config, connectionInfo, jobManagerAkkaURL,
-                      memoryManager, ioManager, network, numberOfSlots) {
+/** Subclass of the [[TaskManager]] to support testing messages
+  *
+  * @param config
+  * @param connectionInfo
+  * @param jobManagerAkkaURL
+  * @param memoryManager
+  * @param ioManager
+  * @param network
+  * @param numberOfSlots
+  */
+class TestingTaskManager(
+    config: TaskManagerConfiguration,
+    connectionInfo: InstanceConnectionInfo,
+    jobManagerAkkaURL: String,
+    memoryManager: DefaultMemoryManager,
+    ioManager: IOManager,
+    network: NetworkEnvironment,
+    numberOfSlots: Int)
+  extends TaskManager(
+    config,
+    connectionInfo,
+    jobManagerAkkaURL,
+    memoryManager,
+    ioManager,
+    network,
+    numberOfSlots) {
 
   import scala.collection.JavaConverters._
 
-
   val waitForRemoval = scala.collection.mutable.HashMap[ExecutionAttemptID, Set[ActorRef]]()
   val waitForJobRemoval = scala.collection.mutable.HashMap[JobID, Set[ActorRef]]()
   val waitForJobManagerToBeTerminated = scala.collection.mutable.HashMap[String, Set[ActorRef]]()
@@ -60,18 +75,19 @@ class TestingTaskManager(config: TaskManagerConfiguration,
 
   var disconnectDisabled = false
 
-
-  override def receiveWithLogMessages = {
-    receiveTestMessages orElse super.receiveWithLogMessages
-  }
-
   /**
    * Handler for testing related messages
    */
-  def receiveTestMessages: Receive = {
+  override def handleMessage: Receive = {
+    handleTestingMessage orElse super.handleMessage
+  }
+
+  def handleTestingMessage: Receive = {
     case NotifyWhenTaskIsRunning(executionID) => {
       Option(runningTasks.get(executionID)) match {
-        case Some(task) if task.getExecutionState == ExecutionState.RUNNING => sender ! true
+        case Some(task) if task.getExecutionState == ExecutionState.RUNNING =>
+          sender ! decorateMessage(true)
+
         case _ =>
           val listeners = waitForRunning.getOrElse(executionID, Set())
           waitForRunning += (executionID -> (listeners + sender))
@@ -79,7 +95,7 @@ class TestingTaskManager(config: TaskManagerConfiguration,
     }
 
     case RequestRunningTasks =>
-      sender ! ResponseRunningTasks(runningTasks.asScala.toMap)
+      sender ! decorateMessage(ResponseRunningTasks(runningTasks.asScala.toMap))
       
     case NotifyWhenTaskRemoved(executionID) =>
       Option(runningTasks.get(executionID)) match {
@@ -88,7 +104,7 @@ class TestingTaskManager(config: TaskManagerConfiguration,
           waitForRemoval += (executionID -> (set + sender))
         case None =>
           if(unregisteredTasks.contains(executionID)) {
-            sender ! true
+            sender ! decorateMessage(true)
           } else {
               val set = waitForRemoval.getOrElse(executionID, Set())
               waitForRemoval += (executionID -> (set + sender))
@@ -96,17 +112,19 @@ class TestingTaskManager(config: TaskManagerConfiguration,
       }
 
     case TaskInFinalState(executionID) =>
-      super.receiveWithLogMessages(TaskInFinalState(executionID))
+      super.handleMessage(TaskInFinalState(executionID))
       waitForRemoval.remove(executionID) match {
-        case Some(actors) => for(actor <- actors) actor ! true
+        case Some(actors) => for(actor <- actors) actor ! decorateMessage(true)
         case None =>
       }
 
       unregisteredTasks += executionID
       
     case RequestBroadcastVariablesWithReferences =>
-      sender ! ResponseBroadcastVariablesWithReferences(
-        bcVarManager.getNumberOfVariablesWithReferences)
+      sender ! decorateMessage(
+        ResponseBroadcastVariablesWithReferences(
+          bcVarManager.getNumberOfVariablesWithReferences)
+      )
 
     case RequestNumActiveConnections =>
       val numActive = if (network.isAssociated) {
@@ -114,30 +132,36 @@ class TestingTaskManager(config: TaskManagerConfiguration,
                       } else {
                         0
                       }
-      sender ! ResponseNumActiveConnections(numActive)
+      sender ! decorateMessage(ResponseNumActiveConnections(numActive))
 
     case NotifyWhenJobRemoved(jobID) =>
       if(runningTasks.values.asScala.exists(_.getJobID == jobID)){
         val set = waitForJobRemoval.getOrElse(jobID, Set())
         waitForJobRemoval += (jobID -> (set + sender))
         import context.dispatcher
-        context.system.scheduler.scheduleOnce(200 milliseconds, this.self, CheckIfJobRemoved(jobID))
+        context.system.scheduler.scheduleOnce(
+          200 milliseconds,
+          this.self,
+          decorateMessage(CheckIfJobRemoved(jobID)))
       }else{
         waitForJobRemoval.get(jobID) match {
-          case Some(listeners) => (listeners + sender) foreach (_ ! true)
-          case None => sender ! true
+          case Some(listeners) => (listeners + sender) foreach (_ ! decorateMessage(true))
+          case None => sender ! decorateMessage(true)
         }
       }
 
     case CheckIfJobRemoved(jobID) =>
       if(runningTasks.values.asScala.forall(_.getJobID != jobID)){
         waitForJobRemoval.remove(jobID) match {
-          case Some(listeners) => listeners foreach (_ ! true)
+          case Some(listeners) => listeners foreach (_ ! decorateMessage(true))
           case None =>
         }
       } else {
         import context.dispatcher
-        context.system.scheduler.scheduleOnce(200 milliseconds, this.self, CheckIfJobRemoved(jobID))
+        context.system.scheduler.scheduleOnce(
+          200 milliseconds,
+          this.self,
+          decorateMessage(CheckIfJobRemoved(jobID)))
       }
 
     case NotifyWhenJobManagerTerminated(jobManager) =>
@@ -160,23 +184,23 @@ class TestingTaskManager(config: TaskManagerConfiguration,
       }
 
     case msg@Terminated(jobManager) =>
-      super.receiveWithLogMessages(msg)
+      super.handleMessage(msg)
 
       waitForJobManagerToBeTerminated.remove(jobManager.path.name) foreach {
         _ foreach {
-          _ ! JobManagerTerminated(jobManager)
+          _ ! decorateMessage(JobManagerTerminated(jobManager))
         }
       }
 
     case msg:Disconnect =>
       if (!disconnectDisabled) {
-        super.receiveWithLogMessages(msg)
+        super.handleMessage(msg)
 
         val jobManager = sender()
 
         waitForJobManagerToBeTerminated.remove(jobManager.path.name) foreach {
           _ foreach {
-            _ ! JobManagerTerminated(jobManager)
+            _ ! decorateMessage(JobManagerTerminated(jobManager))
           }
         }
       }
@@ -185,12 +209,15 @@ class TestingTaskManager(config: TaskManagerConfiguration,
       disconnectDisabled = true
 
     case msg @ UpdateTaskExecutionState(taskExecutionState) =>
-      super.receiveWithLogMessages(msg)
+      super.handleMessage(msg)
 
       if(taskExecutionState.getExecutionState == ExecutionState.RUNNING) {
         waitForRunning.get(taskExecutionState.getID) foreach {
-          _ foreach (_ ! true)
+          _ foreach (_ ! decorateMessage(true))
         }
       }
+
+    case RequestLeaderSessionID =>
+      sender() ! ResponseLeaderSessionID(leaderSessionID)
   }
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-staging/flink-streaming/flink-streaming-connectors/flink-connector-kafka/src/test/resources/log4j-test.properties
----------------------------------------------------------------------
diff --git a/flink-staging/flink-streaming/flink-streaming-connectors/flink-connector-kafka/src/test/resources/log4j-test.properties b/flink-staging/flink-streaming/flink-streaming-connectors/flink-connector-kafka/src/test/resources/log4j-test.properties
index 9ede613..dc20726 100644
--- a/flink-staging/flink-streaming/flink-streaming-connectors/flink-connector-kafka/src/test/resources/log4j-test.properties
+++ b/flink-staging/flink-streaming/flink-streaming-connectors/flink-connector-kafka/src/test/resources/log4j-test.properties
@@ -16,7 +16,7 @@
 # limitations under the License.
 ################################################################################
 
-log4j.rootLogger=INFO, testlogger
+log4j.rootLogger=OFF, testlogger
 
 log4j.appender.testlogger=org.apache.log4j.ConsoleAppender
 log4j.appender.testlogger.target = System.err

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-staging/flink-streaming/flink-streaming-core/src/test/resources/log4j-test.properties
----------------------------------------------------------------------
diff --git a/flink-staging/flink-streaming/flink-streaming-core/src/test/resources/log4j-test.properties b/flink-staging/flink-streaming/flink-streaming-core/src/test/resources/log4j-test.properties
index fcee1b5..0b686e5 100644
--- a/flink-staging/flink-streaming/flink-streaming-core/src/test/resources/log4j-test.properties
+++ b/flink-staging/flink-streaming/flink-streaming-core/src/test/resources/log4j-test.properties
@@ -17,7 +17,7 @@
 ################################################################################
 
 # Set root logger level to DEBUG and its only appender to A1.
-log4j.rootLogger=INFO, A1
+log4j.rootLogger=OFF, A1
 
 # A1 is set to be a ConsoleAppender.
 log4j.appender.A1=org.apache.log4j.ConsoleAppender

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-tests/src/test/java/org/apache/flink/test/accumulators/AccumulatorLiveITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/accumulators/AccumulatorLiveITCase.java b/flink-tests/src/test/java/org/apache/flink/test/accumulators/AccumulatorLiveITCase.java
index f979aea..51d3eb9 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/accumulators/AccumulatorLiveITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/accumulators/AccumulatorLiveITCase.java
@@ -44,6 +44,8 @@ import org.apache.flink.optimizer.plantranslate.JobGraphGenerator;
 import org.apache.flink.runtime.accumulators.AccumulatorRegistry;
 import org.apache.flink.runtime.akka.AkkaUtils;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
+import org.apache.flink.runtime.instance.ActorGateway;
+import org.apache.flink.runtime.instance.AkkaActorGateway;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.runtime.messages.JobManagerMessages;
 import org.apache.flink.runtime.testingUtils.TestingCluster;
@@ -79,7 +81,7 @@ public class AccumulatorLiveITCase {
 	private static final Logger LOG = LoggerFactory.getLogger(AccumulatorLiveITCase.class);
 
 	private static ActorSystem system;
-	private static ActorRef jobManager;
+	private static ActorGateway jobManagerGateway;
 	private static ActorRef taskManager;
 	private static JobID jobID;
 
@@ -104,7 +106,7 @@ public class AccumulatorLiveITCase {
 		config.setString(ConfigConstants.AKKA_ASK_TIMEOUT, TestingUtils.DEFAULT_AKKA_ASK_TIMEOUT());
 		TestingCluster testingCluster = new TestingCluster(config, false, true);
 
-		jobManager = testingCluster.getJobManager();
+		jobManagerGateway = testingCluster.getJobManagerGateway();
 		taskManager = testingCluster.getTaskManagersAsJava().get(0);
 
 		// generate test data
@@ -135,12 +137,15 @@ public class AccumulatorLiveITCase {
 			JobGraph jobGraph = getOptimizedPlan(((PlanExtractor) env).plan);
 			jobID = jobGraph.getJobID();
 
+			ActorGateway selfGateway = new AkkaActorGateway(getRef(), jobManagerGateway.leaderSessionID());
+
 			// register for accumulator changes
-			jobManager.tell(new TestingJobManagerMessages.NotifyWhenAccumulatorChange(jobID), getRef());
+			jobManagerGateway.tell(new TestingJobManagerMessages.NotifyWhenAccumulatorChange(jobID), selfGateway);
 			expectMsgEquals(TIMEOUT, true);
 
 			// submit job
-			jobManager.tell(new JobManagerMessages.SubmitJob(jobGraph, false), getRef());
+
+			jobManagerGateway.tell(new JobManagerMessages.SubmitJob(jobGraph, false), selfGateway);
 			expectMsgClass(TIMEOUT, Status.Success.class);
 
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-tests/src/test/java/org/apache/flink/test/cancelling/CancellingTestBase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/cancelling/CancellingTestBase.java b/flink-tests/src/test/java/org/apache/flink/test/cancelling/CancellingTestBase.java
index a3186a7..1502110 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/cancelling/CancellingTestBase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/cancelling/CancellingTestBase.java
@@ -21,7 +21,6 @@ package org.apache.flink.test.cancelling;
 
 import java.util.concurrent.atomic.AtomicReference;
 
-import akka.actor.ActorRef;
 import org.apache.flink.configuration.ConfigConstants;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.client.JobCancellationException;
@@ -115,7 +114,7 @@ public abstract class CancellingTestBase {
 				public void run() {
 					try {
 						Thread.sleep(msecsTillCanceling);
-						executor.getJobManager().tell(new JobManagerMessages.CancelJob(jobGraph.getJobID()), ActorRef.noSender());
+						executor.getJobManagerGateway().tell(new JobManagerMessages.CancelJob(jobGraph.getJobID()));
 					}
 					catch (Throwable t) {
 						error.set(t);

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-tests/src/test/java/org/apache/flink/test/runtime/minicluster/LocalFlinkMiniClusterITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/runtime/minicluster/LocalFlinkMiniClusterITCase.java b/flink-tests/src/test/java/org/apache/flink/test/runtime/minicluster/LocalFlinkMiniClusterITCase.java
index 7933932..c45cfe8 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/runtime/minicluster/LocalFlinkMiniClusterITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/runtime/minicluster/LocalFlinkMiniClusterITCase.java
@@ -18,18 +18,22 @@
 
 package org.apache.flink.test.runtime.minicluster;
 
-import akka.actor.ActorRef;
 import akka.actor.ActorSystem;
 import akka.testkit.JavaTestKit;
 import org.apache.flink.configuration.ConfigConstants;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.akka.AkkaUtils;
+import org.apache.flink.runtime.instance.AkkaActorGateway;
+import org.apache.flink.runtime.instance.ActorGateway;
 import org.apache.flink.runtime.messages.JobManagerMessages;
 import org.apache.flink.runtime.minicluster.LocalFlinkMiniCluster;
 import org.apache.flink.runtime.testingUtils.TestingUtils;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import scala.Option;
+
+import java.util.UUID;
 
 public class LocalFlinkMiniClusterITCase {
 
@@ -59,19 +63,24 @@ public class LocalFlinkMiniClusterITCase {
 			config.setInteger(ConfigConstants.TASK_MANAGER_NUM_TASK_SLOTS, numSlots);
 			miniCluster = new LocalFlinkMiniCluster(config, true);
 
-			final ActorRef jm = miniCluster.getJobManager();
+			final ActorGateway jmGateway = miniCluster.getJobManagerGateway();
 
 			new JavaTestKit(system) {{
+				final ActorGateway selfGateway = new AkkaActorGateway(getRef(), Option.<UUID>empty());
+
 				new Within(TestingUtils.TESTING_DURATION()) {
 
 					@Override
 					protected void run() {
-						jm.tell(JobManagerMessages.getRequestNumberRegisteredTaskManager(),
-								getRef());
+						jmGateway.tell(
+								JobManagerMessages.getRequestNumberRegisteredTaskManager(),
+								selfGateway);
 
 						expectMsgEquals(TestingUtils.TESTING_DURATION(), numTMs);
 
-						jm.tell(JobManagerMessages.getRequestTotalNumberOfSlots(), getRef());
+						jmGateway.tell(
+								JobManagerMessages.getRequestTotalNumberOfSlots(),
+								selfGateway);
 
 						expectMsgEquals(TestingUtils.TESTING_DURATION(), numTMs*numSlots);
 					}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/jobmanager/JobManagerFailsITCase.scala
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/jobmanager/JobManagerFailsITCase.scala b/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/jobmanager/JobManagerFailsITCase.scala
index a590ae7..feb90b2 100644
--- a/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/jobmanager/JobManagerFailsITCase.scala
+++ b/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/jobmanager/JobManagerFailsITCase.scala
@@ -32,12 +32,17 @@ import org.apache.flink.runtime.jobmanager.Tasks.{BlockingNoOpInvokable, NoOpInv
 import org.apache.flink.runtime.messages.JobManagerMessages._
 import org.apache.flink.runtime.testingUtils.TestingMessages.DisableDisconnect
 import org.apache.flink.runtime.testingUtils.TestingTaskManagerMessages.{JobManagerTerminated, NotifyWhenJobManagerTerminated}
-import org.apache.flink.runtime.testingUtils.TestingUtils
+import org.apache.flink.runtime.testingUtils.{ScalaTestingUtils, TestingUtils}
 import org.apache.flink.test.util.ForkableFlinkMiniCluster
 
 @RunWith(classOf[JUnitRunner])
-class JobManagerFailsITCase(_system: ActorSystem) extends TestKit(_system) with ImplicitSender
-with WordSpecLike with Matchers with BeforeAndAfterAll {
+class JobManagerFailsITCase(_system: ActorSystem)
+  extends TestKit(_system)
+  with ImplicitSender
+  with WordSpecLike
+  with Matchers
+  with BeforeAndAfterAll
+  with ScalaTestingUtils {
 
   def this() = this(ActorSystem("TestingActorSystem", AkkaUtils.getDefaultAkkaConfig))
 
@@ -52,18 +57,18 @@ with WordSpecLike with Matchers with BeforeAndAfterAll {
       val cluster = startDeathwatchCluster(num_slots, 1)
 
       val tm = cluster.getTaskManagers(0)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       // disable disconnect message to test death watch
       tm ! DisableDisconnect
 
       try{
-        jm ! RequestNumberRegisteredTaskManager
+        jmGateway.tell(RequestNumberRegisteredTaskManager, self)
         expectMsg(1)
 
-        tm ! NotifyWhenJobManagerTerminated(jm)
+        tm ! NotifyWhenJobManagerTerminated(jmGateway.actor)
 
-        jm ! PoisonPill
+        jmGateway.tell(PoisonPill, self)
 
         expectMsgClass(classOf[JobManagerTerminated])
 
@@ -71,7 +76,7 @@ with WordSpecLike with Matchers with BeforeAndAfterAll {
 
         cluster.waitForTaskManagersToBeRegistered()
 
-        cluster.getJobManager ! RequestNumberRegisteredTaskManager
+        cluster.getJobManagerGateway.tell(RequestNumberRegisteredTaskManager, self)
 
         expectMsg(1)
       } finally {
@@ -94,27 +99,27 @@ with WordSpecLike with Matchers with BeforeAndAfterAll {
 
       val cluster = startDeathwatchCluster(num_slots / 2, 2)
 
-      var jm = cluster.getJobManager
+      var jmGateway = cluster.getJobManagerGateway
       val tm = cluster.getTaskManagers(0)
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
 
-          tm ! NotifyWhenJobManagerTerminated(jm)
+          tm.tell(NotifyWhenJobManagerTerminated(jmGateway.actor()), self)
 
-          jm ! PoisonPill
+          jmGateway.tell(PoisonPill, self)
 
           expectMsgClass(classOf[JobManagerTerminated])
 
           cluster.restartJobManager()
 
-          jm = cluster.getJobManager
+          jmGateway = cluster.getJobManagerGateway
 
           cluster.waitForTaskManagersToBeRegistered()
 
-          jm ! SubmitJob(jobGraph2, false)
+          jmGateway.tell(SubmitJob(jobGraph2, false), self)
 
           val failure = expectMsgType[Success]
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/jobmanager/JobManagerLeaderSessionIDITSuite.scala
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/jobmanager/JobManagerLeaderSessionIDITSuite.scala b/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/jobmanager/JobManagerLeaderSessionIDITSuite.scala
new file mode 100644
index 0000000..010a086
--- /dev/null
+++ b/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/jobmanager/JobManagerLeaderSessionIDITSuite.scala
@@ -0,0 +1,112 @@
+/*
+ * 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.api.scala.runtime.jobmanager
+
+import java.util.UUID
+
+import akka.actor.ActorSystem
+import akka.actor.Status.Success
+import akka.testkit.{ImplicitSender, TestKit}
+import org.apache.flink.runtime.akka.AkkaUtils
+import org.apache.flink.runtime.jobgraph.tasks.AbstractInvokable
+import org.apache.flink.runtime.jobgraph.{JobGraph, JobVertex}
+import org.apache.flink.runtime.messages.JobManagerMessages.{LeaderSessionMessage, CancelJob,
+JobResultSuccess, SubmitJob}
+import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages.{AllVerticesRunning,
+WaitForAllVerticesToBeRunning}
+import org.apache.flink.runtime.testingUtils.{ScalaTestingUtils, TestingUtils}
+import org.junit.runner.RunWith
+import org.scalatest.{FunSuiteLike, Matchers, BeforeAndAfterAll}
+import org.scalatest.junit.JUnitRunner
+
+@RunWith(classOf[JUnitRunner])
+class JobManagerLeaderSessionIDITSuite(_system: ActorSystem)
+  extends TestKit(_system)
+  with ImplicitSender
+  with FunSuiteLike
+  with Matchers
+  with BeforeAndAfterAll
+  with ScalaTestingUtils {
+
+  val numTaskManagers = 2
+  val taskManagerNumSlots = 2
+  val numSlots = numTaskManagers * taskManagerNumSlots
+
+  val cluster = TestingUtils.startTestingCluster(
+    taskManagerNumSlots,
+    numTaskManagers,
+    TestingUtils.DEFAULT_AKKA_ASK_TIMEOUT);
+
+  def this() = this(ActorSystem("TestingActorSystem", AkkaUtils.getDefaultAkkaConfig))
+
+  override def afterAll(): Unit = {
+    cluster.stop()
+    TestKit.shutdownActorSystem(system)
+  }
+
+  test("A JobManager should not process CancelJob messages with the wrong leader session ID") {
+    val sender = new JobVertex("BlockingSender");
+    sender.setParallelism(numSlots)
+    sender.setInvokableClass(classOf[BlockingUntilSignalNoOpInvokable])
+    val jobGraph = new JobGraph("TestJob", sender)
+
+    val oldSessionID = Option(UUID.randomUUID())
+
+    val jmGateway = cluster.getJobManagerGateway()
+    val jm = jmGateway.actor()
+
+    within(TestingUtils.TESTING_DURATION) {
+      jmGateway.tell(SubmitJob(jobGraph, false), self)
+
+      expectMsg(Success(jobGraph.getJobID))
+
+      jmGateway.tell(WaitForAllVerticesToBeRunning(jobGraph.getJobID), self)
+
+      expectMsg(AllVerticesRunning(jobGraph.getJobID))
+
+      jm ! LeaderSessionMessage(oldSessionID, CancelJob(jobGraph.getJobID))
+
+      BlockingUntilSignalNoOpInvokable.triggerExecution()
+
+      expectMsgType[JobResultSuccess]
+    }
+  }
+}
+
+class BlockingUntilSignalNoOpInvokable extends AbstractInvokable {
+  override def registerInputOutput(): Unit = {
+
+  }
+
+  override def invoke(): Unit = {
+    BlockingUntilSignalNoOpInvokable.lock.synchronized{
+      BlockingUntilSignalNoOpInvokable.lock.wait()
+    }
+  }
+}
+
+object BlockingUntilSignalNoOpInvokable {
+  val lock = new Object
+
+  def triggerExecution(): Unit = {
+    lock.synchronized{
+      lock.notifyAll()
+    }
+  }
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/taskmanager/TaskManagerFailsITCase.scala
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/taskmanager/TaskManagerFailsITCase.scala b/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/taskmanager/TaskManagerFailsITCase.scala
index c20f621..d108798 100644
--- a/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/taskmanager/TaskManagerFailsITCase.scala
+++ b/flink-tests/src/test/scala/org/apache/flink/api/scala/runtime/taskmanager/TaskManagerFailsITCase.scala
@@ -32,7 +32,7 @@ import org.apache.flink.runtime.messages.JobManagerMessages.{JobResultSuccess, R
 import org.apache.flink.runtime.messages.TaskManagerMessages.{RegisteredAtJobManager, NotifyWhenRegisteredAtJobManager}
 import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages._
 import org.apache.flink.runtime.testingUtils.TestingMessages.DisableDisconnect
-import org.apache.flink.runtime.testingUtils.TestingUtils
+import org.apache.flink.runtime.testingUtils.{ScalaTestingUtils, TestingUtils}
 import org.apache.flink.test.util.ForkableFlinkMiniCluster
 
 import org.junit.runner.RunWith
@@ -40,8 +40,13 @@ import org.scalatest.junit.JUnitRunner
 import org.scalatest.{BeforeAndAfterAll, Matchers, WordSpecLike}
 
 @RunWith(classOf[JUnitRunner])
-class TaskManagerFailsITCase(_system: ActorSystem) extends TestKit(_system) with ImplicitSender
-with WordSpecLike with Matchers with BeforeAndAfterAll {
+class TaskManagerFailsITCase(_system: ActorSystem)
+  extends TestKit(_system)
+  with ImplicitSender
+  with WordSpecLike
+  with Matchers
+  with BeforeAndAfterAll
+  with ScalaTestingUtils {
 
   def this() = this(ActorSystem("TestingActorSystem", AkkaUtils.getDefaultAkkaConfig))
 
@@ -57,22 +62,22 @@ with WordSpecLike with Matchers with BeforeAndAfterAll {
       val cluster = startDeathwatchCluster(num_slots, 2)
 
       val taskManagers = cluster.getTaskManagers
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
-      jm ! DisableDisconnect
+      jmGateway.tell(DisableDisconnect)
 
       try{
         within(TestingUtils.TESTING_DURATION){
-          jm ! RequestNumberRegisteredTaskManager
+          jmGateway.tell(RequestNumberRegisteredTaskManager, self)
           expectMsg(2)
 
-          jm ! NotifyWhenTaskManagerTerminated(taskManagers(0))
+          jmGateway.tell(NotifyWhenTaskManagerTerminated(taskManagers(0)), self)
 
           taskManagers(0) ! PoisonPill
 
           val TaskManagerTerminated(tm) = expectMsgClass(classOf[TaskManagerTerminated])
 
-          jm ! RequestNumberRegisteredTaskManager
+          jmGateway.tell(RequestNumberRegisteredTaskManager, self)
           expectMsg(1)
         }
       }
@@ -97,18 +102,18 @@ with WordSpecLike with Matchers with BeforeAndAfterAll {
 
       val cluster = ForkableFlinkMiniCluster.startCluster(num_tasks, 2)
 
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
 
-          jm ! WaitForAllVerticesToBeRunningOrFinished(jobID)
+          jmGateway.tell(WaitForAllVerticesToBeRunningOrFinished(jobID), self)
 
           expectMsg(AllVerticesRunning(jobID))
 
-          jm ! RequestWorkingTaskManager(jobID)
+          jmGateway.tell(RequestWorkingTaskManager(jobID), self)
 
           val gatewayOption = expectMsgType[WorkingTaskManager].gatewayOption
 
@@ -150,14 +155,14 @@ with WordSpecLike with Matchers with BeforeAndAfterAll {
       val cluster = ForkableFlinkMiniCluster.startCluster(num_tasks, 2)
 
       val taskManagers = cluster.getTaskManagers
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try {
         within(TestingUtils.TESTING_DURATION) {
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
 
-          jm ! WaitForAllVerticesToBeRunningOrFinished(jobID)
+          jmGateway.tell(WaitForAllVerticesToBeRunningOrFinished(jobID), self)
           expectMsg(AllVerticesRunning(jobID))
 
           // kill one task manager
@@ -193,11 +198,11 @@ with WordSpecLike with Matchers with BeforeAndAfterAll {
       val cluster = startDeathwatchCluster(num_slots/2, 2)
 
       var tm = cluster.getTaskManagers(0)
-      val jm = cluster.getJobManager
+      val jmGateway = cluster.getJobManagerGateway
 
       try{
         within(TestingUtils.TESTING_DURATION){
-          jm ! SubmitJob(jobGraph, false)
+          jmGateway.tell(SubmitJob(jobGraph, false), self)
           expectMsg(Success(jobGraph.getJobID))
 
           tm ! PoisonPill
@@ -219,7 +224,7 @@ with WordSpecLike with Matchers with BeforeAndAfterAll {
 
           expectMsg(RegisteredAtJobManager)
 
-          jm ! SubmitJob(jobGraph2, false)
+          jmGateway.tell(SubmitJob(jobGraph2, false), self)
 
           expectMsgType[Success]
 

http://git-wip-us.apache.org/repos/asf/flink/blob/22224542/flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationClient.scala
----------------------------------------------------------------------
diff --git a/flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationClient.scala b/flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationClient.scala
index ec980d0..6930be6 100644
--- a/flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationClient.scala
+++ b/flink-yarn/src/main/scala/org/apache/flink/yarn/ApplicationClient.scala
@@ -22,8 +22,9 @@ import java.net.InetSocketAddress
 
 import akka.actor._
 import akka.pattern.ask
+import grizzled.slf4j.Logger
 import org.apache.flink.configuration.Configuration
-import org.apache.flink.runtime.{ActorSynchronousLogging, ActorLogMessages}
+import org.apache.flink.runtime.{FlinkActor, LogMessages}
 import org.apache.flink.runtime.akka.AkkaUtils
 import org.apache.flink.runtime.jobmanager.JobManager
 import org.apache.flink.runtime.yarn.FlinkYarnClusterStatus
@@ -35,11 +36,11 @@ import scala.language.postfixOps
 import scala.util.{Failure, Success}
 
 class ApplicationClient(flinkConfig: Configuration)
-  extends Actor
-  with ActorLogMessages
-  with ActorSynchronousLogging {
+  extends FlinkActor with LogMessages {
   import context._
 
+  val log = Logger(getClass)
+
   val INITIAL_POLLING_DELAY = 0 seconds
   val WAIT_FOR_YARN_INTERVAL = 2 seconds
   val POLLING_INTERVAL = 3 seconds
@@ -70,7 +71,7 @@ class ApplicationClient(flinkConfig: Configuration)
     context.system.shutdown()
   }
 
-  override def receiveWithLogMessages: Receive = {
+  override def handleMessage: Receive = {
     // ----------------------------- Registration -> Status updates -> shutdown ----------------
     case LocalRegisterClient(address: InetSocketAddress) =>
       val jmAkkaUrl = JobManager.getRemoteJobManagerAkkaURL(address)
@@ -78,13 +79,13 @@ class ApplicationClient(flinkConfig: Configuration)
       val jobManagerFuture = AkkaUtils.getReference(jmAkkaUrl, system, timeout)
 
       jobManagerFuture.onComplete {
-        case Success(jm) => self ! JobManagerActorRef(jm)
+        case Success(jm) => self ! decorateMessage(JobManagerActorRef(jm))
         case Failure(t) =>
           log.error("Registration at JobManager/ApplicationMaster failed. Shutting " +
             "ApplicationClient down.", t)
 
           // we could not connect to the job manager --> poison ourselves
-          self ! PoisonPill
+          self ! decorateMessage(PoisonPill)
       }
 
     case JobManagerActorRef(jm) =>
@@ -93,40 +94,45 @@ class ApplicationClient(flinkConfig: Configuration)
       // the message came from the FlinkYarnCluster. We send the message to the JobManager.
       // it is important not to forward the message because the JobManager is storing the
       // sender as the Application Client (this class).
-      (jm ? RegisterClient(self))(timeout).onFailure{
+      (jm ? decorateMessage(RegisterClient(self)))(timeout).onFailure{
         case t: Throwable =>
           log.error("Could not register at the job manager.", t)
-          self ! PoisonPill
+          self ! decorateMessage(PoisonPill)
       }
 
       // schedule a periodic status report from the JobManager
       // request the number of task managers and slots from the job manager
-      pollingTimer = Some(context.system.scheduler.schedule(INITIAL_POLLING_DELAY,
-        WAIT_FOR_YARN_INTERVAL, jm, PollYarnClusterStatus))
+      pollingTimer = Some(
+        context.system.scheduler.schedule(
+          INITIAL_POLLING_DELAY,
+          WAIT_FOR_YARN_INTERVAL,
+          jm,
+          PollYarnClusterStatus)
+      )
 
     case LocalUnregisterClient =>
       // unregister client from AM
       yarnJobManager foreach {
-        _ ! UnregisterClient
+        _ ! decorateMessage(UnregisterClient)
       }
       // poison ourselves
-      self ! PoisonPill
+      self ! decorateMessage(PoisonPill)
 
     case msg: StopYarnSession =>
       log.info("Sending StopYarnSession request to ApplicationMaster.")
       stopMessageReceiver = Some(sender)
       yarnJobManager foreach {
-        _ forward msg
+        _ forward decorateMessage(msg)
       }
 
     case JobManagerStopped =>
       log.info("Remote JobManager has been stopped successfully. " +
         "Stopping local application client")
       stopMessageReceiver foreach {
-        _ ! JobManagerStopped
+        _ ! decorateMessage(JobManagerStopped)
       }
       // poison ourselves
-      self ! PoisonPill
+      self ! decorateMessage(PoisonPill)
 
     // handle the responses from the PollYarnClusterStatus messages to the yarn job mgr
     case status: FlinkYarnClusterStatus =>
@@ -135,12 +141,12 @@ class ApplicationClient(flinkConfig: Configuration)
 
     // locally get cluster status
     case LocalGetYarnClusterStatus =>
-      sender() ! latestClusterStatus
+      sender() ! decorateMessage(latestClusterStatus)
 
     // Forward message to Application Master
     case msg: StopAMAfterJob =>
       yarnJobManager foreach {
-        _ forward msg
+        _ forward decorateMessage(msg)
       }
 
     // -----------------  handle messages from the cluster -------------------
@@ -152,9 +158,9 @@ class ApplicationClient(flinkConfig: Configuration)
     // locally forward messages
     case LocalGetYarnMessage =>
       if(messagesQueue.size > 0) {
-        sender() ! Option(messagesQueue.dequeue)
+        sender() ! decorateMessage(Option(messagesQueue.dequeue))
       } else {
-        sender() ! None
+        sender() ! decorateMessage(None)
       }
   }