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,
      byte[] secretShuffleKey); @Override InetSocketAddress getAddress(); @Override boolean canCommit(TaskAttemptID taskAttemptID); @Override void commitPending(TaskAttemptID taskAttemptID, TaskStatus taskStatsu); @Override void done(TaskAttemptID taskAttemptID); @Override void fatalError(TaskAttemptID taskAttemptID, String msg); @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 boolean ping(TaskAttemptID taskAttemptID); @Override void reportDiagnosticInfo(TaskAttemptID taskAttemptID, String diagnosticInfo); @Override boolean 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);  }
@Test (timeout=5000) public void testGetTask() throws IOException { AppContext appCtx = mock(AppContext.class); JobTokenSecretManager secret = mock(JobTokenSecretManager.class); RMHeartbeatHandler rmHeartbeatHandler = mock(RMHeartbeatHandler.class); TaskHeartbeatHandler hbHandler = mock(TaskHeartbeatHandler.class); MockTaskAttemptListenerImpl listener = new MockTaskAttemptListenerImpl(appCtx, secret, rmHeartbeatHandler, hbHandler); Configuration conf = new Configuration(); listener.init(conf); listener.start(); JVMId id = new JVMId("foo",1, true, 1); WrappedJvmID wid = new WrappedJvmID(id.getJobId(), id.isMap, id.getId()); JvmContext context = new JvmContext(); context.jvmId = id; JvmTask result = listener.getTask(context); assertNotNull(result); assertTrue(result.shouldDie); TaskAttemptId attemptID = mock(TaskAttemptId.class); Task task = mock(Task.class); 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); listener.stop(); JVMId jvmid = JVMId.forName("jvm_001_002_m_004"); assertNotNull(jvmid); try { JVMId.forName("jvm_001_002_m_004_006"); Assert.fail(); } catch (IllegalArgumentException e) { assertEquals(e.getMessage(), "TaskId string : jvm_001_002_m_004_006 is not properly formed"); } }