TaskAttemptListenerImpl extends CompositeService implements TaskUmbilicalProtocol, TaskAttemptListener { @Override public JvmTask getTask(JvmContext context) throws IOException { JVMId jvmId = context.jvmId; LOG.info("JVM with ID : " + jvmId + " asked for a task"); JvmTask jvmTask = null; WrappedJvmID wJvmID = new WrappedJvmID(jvmId.getJobId(), jvmId.isMap, jvmId.getId()); if (!jvmIDToActiveAttemptMap.containsKey(wJvmID)) { LOG.info("JVM with ID: " + jvmId + " is invalid and will be killed."); jvmTask = TASK_FOR_INVALID_JVM; } else { if (!launchedJVMs.contains(wJvmID)) { jvmTask = null; LOG.info("JVM with ID: " + jvmId + " asking for task before AM launch registered. Given null task"); } else { org.apache.hadoop.mapred.Task task = jvmIDToActiveAttemptMap.remove(wJvmID); launchedJVMs.remove(wJvmID); LOG.info("JVM with ID: " + jvmId + " given task: " + task.getTaskID()); task.setEncryptedSpillKey(encryptedSpillKey); jvmTask = new JvmTask(task, false); } } return jvmTask; } TaskAttemptListenerImpl(AppContext context,
      JobTokenSecretManager jobTokenSecretManager,
      RMHeartbeatHandler rmHeartbeatHandler,
      AMPreemptionPolicy preemptionPolicy);  TaskAttemptListenerImpl(AppContext context,
      JobTokenSecretManager jobTokenSecretManager,
      RMHeartbeatHandler rmHeartbeatHandler,
      AMPreemptionPolicy preemptionPolicy, byte[] secretShuffleKey); @Override InetSocketAddress getAddress(); @Override boolean canCommit(TaskAttemptID taskAttemptID); @Override void commitPending(TaskAttemptID taskAttemptID, TaskStatus taskStatsu); @Override void preempted(TaskAttemptID taskAttemptID, TaskStatus taskStatus); @Override void done(TaskAttemptID taskAttemptID); @Override void fatalError(TaskAttemptID taskAttemptID, String msg, boolean fastFail); @Override void fsError(TaskAttemptID taskAttemptID, String message); @Override void shuffleError(TaskAttemptID taskAttemptID, String message); @Override MapTaskCompletionEventsUpdate getMapCompletionEvents(
      JobID jobIdentifier, int startIndex, int maxEvents,
      TaskAttemptID taskAttemptID); @Override void reportDiagnosticInfo(TaskAttemptID taskAttemptID, String diagnosticInfo); @Override AMFeedback statusUpdate(TaskAttemptID taskAttemptID,
      TaskStatus taskStatus); @Override long getProtocolVersion(String arg0, long arg1); @Override void reportNextRecordRange(TaskAttemptID taskAttemptID, Range range); @Override JvmTask getTask(JvmContext context); @Override void registerPendingTask(
      org.apache.hadoop.mapred.Task task, WrappedJvmID jvmID); @Override void registerLaunchedTask(
      org.apache.hadoop.mapreduce.v2.api.records.TaskAttemptId attemptID,
      WrappedJvmID jvmId); @Override void unregister(
      org.apache.hadoop.mapreduce.v2.api.records.TaskAttemptId attemptID,
      WrappedJvmID jvmID); @Override ProtocolSignature getProtocolSignature(String protocol,
      long clientVersion, int clientMethodsHash); @Override TaskCheckpointID getCheckpointID(TaskID taskId); @Override void setCheckpointID(TaskID taskId, TaskCheckpointID cid);  }
@Test (timeout=5000) public void testGetTask() throws IOException { configureMocks(); startListener(false); JvmContext context = new JvmContext(); context.jvmId = id; JvmTask result = listener.getTask(context); assertNotNull(result); assertTrue(result.shouldDie); listener.registerPendingTask(task, wid); result = listener.getTask(context); assertNull(result); listener.unregister(attemptId, wid); listener.registerPendingTask(task, wid); listener.registerLaunchedTask(attemptId, wid); verify(hbHandler).register(attemptId); result = listener.getTask(context); assertNotNull(result); assertFalse(result.shouldDie); result = listener.getTask(context); assertNotNull(result); assertTrue(result.shouldDie); listener.unregister(attemptId, wid); result = listener.getTask(context); assertNotNull(result); assertTrue(result.shouldDie); JVMId jvmid = JVMId.forName("jvm_001_002_m_004"); assertNotNull(jvmid); try { JVMId.forName("jvm_001_002_m_004_006"); fail(); } catch (IllegalArgumentException e) { assertThat(e.getMessage()).isEqualTo( "TaskId string : jvm_001_002_m_004_006 is not properly formed"); } }