Java Code Examples for org.apache.tez.common.TezUtils#createUserPayloadFromConf()

static OutputContext createOutputContext() throws IOException {
  OutputContext outputContext = mock(OutputContext.class);
  Configuration conf = new TezConfiguration();
  UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
  String[] workingDirs = new String[]{"workDir1"};
  OutputStatisticsReporter statsReporter = mock(OutputStatisticsReporter.class);
  TezCounters counters = new TezCounters();

  doReturn("destinationVertex").when(outputContext).getDestinationVertexName();
  doReturn(payLoad).when(outputContext).getUserPayload();
  doReturn(workingDirs).when(outputContext).getWorkDirs();
  doReturn(200 * 1024 * 1024l).when(outputContext).getTotalMemoryAvailableToTask();
  doReturn(counters).when(outputContext).getCounters();
  doReturn(statsReporter).when(outputContext).getStatisticsReporter();
  doReturn(new Configuration(false)).when(outputContext).getContainerConfiguration();
  return outputContext;
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  byte[] payload = null;
  int taskCount = TEZ_SIMPLE_DAG_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(TEZ_SIMPLE_DAG_NUM_TASKS, TEZ_SIMPLE_DAG_NUM_TASKS_DEFAULT);
    payload = TezUtils.createUserPayloadFromConf(conf);
  }
  DAG dag = new DAG(name);
  Vertex v1 = new Vertex("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = new Vertex("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v3 = new Vertex("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addEdge(new Edge(v1, v2, 
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, 
          SchedulingType.SEQUENTIAL, 
          TestOutput.getOutputDesc(payload), 
          TestInput.getInputDesc(payload))));
  dag.addVertex(v3).addEdge(new Edge(v2, v3, 
          new EdgeProperty(DataMovementType.SCATTER_GATHER, 
              DataSourceType.PERSISTED, 
              SchedulingType.SEQUENTIAL, 
              TestOutput.getOutputDesc(payload), 
              TestInput.getInputDesc(payload))));
  return dag;
}
 

private void _testPipelinedShuffle(String sorterName) throws Exception {
  conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 3);

  conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_SORTER_CLASS, sorterName);
  conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, false);
  conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_PIPELINED_SHUFFLE_ENABLED, true);
  OutputContext context = createTezOutputContext();
  UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
  doReturn(payLoad).when(context).getUserPayload();
  sortedOutput = new OrderedPartitionedKVOutput(context, partitions);

  sortedOutput.initialize();
  sortedOutput.start();

  assertFalse(sortedOutput.finalMergeEnabled);
  assertTrue(sortedOutput.pipelinedShuffle);
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  UserPayload payload = UserPayload.create(null);
  int taskCount = TEZ_SIMPLE_DAG_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(TEZ_SIMPLE_DAG_NUM_TASKS, TEZ_SIMPLE_DAG_NUM_TASKS_DEFAULT);
    payload = TezUtils.createUserPayloadFromConf(conf);
  }
  DAG dag = DAG.create(name);
  Vertex v1 = Vertex.create("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = Vertex.create("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addEdge(Edge.create(v1, v2,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  return dag;
}
 

public static DAG createDAG(String name, 
        Configuration conf) throws Exception {
    if (conf != null) {
      payload = TezUtils.createUserPayloadFromConf(conf);
    } 
    dag = DAG.create(name);
    addDAGVerticesAndEdges();
    return dag;
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  UserPayload payload = UserPayload.create(null);
  int taskCount = TEZ_SIMPLE_V_DAG_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(TEZ_SIMPLE_V_DAG_NUM_TASKS, TEZ_SIMPLE_V_DAG_NUM_TASKS_DEFAULT);
    payload = TezUtils.createUserPayloadFromConf(conf);
  }
  DAG dag = DAG.create(name);
  Vertex v1 = Vertex.create("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = Vertex.create("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v3 = Vertex.create("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addVertex(v3);
  dag.addEdge(Edge.create(v1, v3,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  dag.addEdge(Edge.create(v2, v3,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  return dag;
}
 

/**
 * Creates the user payload to be set on the OutputDescriptor for MROutput
 * @param conf Configuration for the OutputFormat
 * @param outputFormatName Name of the class of the OutputFormat
 * @param useNewApi Use new mapreduce API or old mapred API
 * @return
 * @throws IOException
 */
public static byte[] createUserPayload(Configuration conf, 
    String outputFormatName, boolean useNewApi) throws IOException {
  Configuration outputConf = new JobConf(conf);
  outputConf.set(MRJobConfig.OUTPUT_FORMAT_CLASS_ATTR, outputFormatName);
  outputConf.setBoolean("mapred.mapper.new-api", useNewApi);
  MRHelpers.translateVertexConfToTez(outputConf);
  MRHelpers.doJobClientMagic(outputConf);
  return TezUtils.createUserPayloadFromConf(outputConf);
}
 

/**
 * Get a byte array representation of the configuration
 * @return a byte array which can be used as the payload
 */
public byte[] toByteArray() {
  try {
    return TezUtils.createUserPayloadFromConf(conf);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  UserPayload payload = UserPayload.create(null);
  int taskCount = TEZ_SIMPLE_REVERSE_V_DAG_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(TEZ_SIMPLE_REVERSE_V_DAG_NUM_TASKS, TEZ_SIMPLE_REVERSE_V_DAG_NUM_TASKS_DEFAULT);
    payload = TezUtils.createUserPayloadFromConf(conf);
  }
  DAG dag = DAG.create(name);
  Vertex v1 = Vertex.create("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = Vertex.create("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v3 = Vertex.create("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addVertex(v3);
  dag.addEdge(Edge.create(v1, v2,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  dag.addEdge(Edge.create(v1, v3,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  return dag;
}
 

/**
 * Get a UserPayload representation of the Configuration
 * @return a {@link org.apache.tez.dag.api.UserPayload} instance
 */
public UserPayload toUserPayload() {
  try {
    return TezUtils.createUserPayloadFromConf(conf);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
}
 

public static DAG createDAG(String name, 
    Configuration conf) throws Exception {
  byte[] payload = null;
  int taskCount = TEZ_SIMPLE_V_DAG_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(TEZ_SIMPLE_V_DAG_NUM_TASKS, TEZ_SIMPLE_V_DAG_NUM_TASKS_DEFAULT);
    payload = TezUtils.createUserPayloadFromConf(conf);
  }
  DAG dag = new DAG(name);
  Vertex v1 = new Vertex("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v2 = new Vertex("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v3 = new Vertex("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  dag.addVertex(v1).addVertex(v2).addVertex(v3);
  dag.addEdge(new Edge(v1, v3, 
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, 
          SchedulingType.SEQUENTIAL, 
          TestOutput.getOutputDesc(payload), 
          TestInput.getInputDesc(payload))));
  dag.addEdge(new Edge(v2, v3, 
          new EdgeProperty(DataMovementType.SCATTER_GATHER, 
              DataSourceType.PERSISTED, 
              SchedulingType.SEQUENTIAL, 
              TestOutput.getOutputDesc(payload), 
              TestInput.getInputDesc(payload))));
  return dag;
}
 

public static DAG createDAG(String name, 
        Configuration conf) throws Exception {
    if (conf != null) {
      payload = TezUtils.createUserPayloadFromConf(conf);
    } 
    dag = new DAG(name);
    addDAGVerticesAndEdges();
    return dag;
}
 

public static DAG createDAG(String name, 
        Configuration conf) throws Exception {
    if (conf != null) {
      payload = TezUtils.createUserPayloadFromConf(conf);
    } 
    dag = new DAG(name);
    addDAGVerticesAndEdges();
    return dag;
}
 

public static DAG createDAG(String name, 
        Configuration conf) throws Exception {
    if (conf != null) {
      payload = TezUtils.createUserPayloadFromConf(conf);
    } 
    dag = new DAG(name);
    addDAGVerticesAndEdges();
    return dag;
}
 

private void startSortedOutput(int partitions) throws Exception {
  OutputContext context = createTezOutputContext();
  conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, 4);
  UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
  doReturn(payLoad).when(context).getUserPayload();
  sortedOutput = new OrderedPartitionedKVOutput(context, partitions);
  sortedOutput.initialize();
  sortedOutput.start();
  writer = sortedOutput.getWriter();
}
 

@Test(timeout = 5000)
public void testMultipleContainerLaunchers() throws IOException, TezException {
  Configuration conf = new Configuration(false);
  conf.set("testkey", "testvalue");
  UserPayload userPayload = TezUtils.createUserPayloadFromConf(conf);

  AppContext appContext = mock(AppContext.class);
  TaskCommunicatorManagerInterface tal = mock(TaskCommunicatorManagerInterface.class);

  String customLauncherName = "customLauncher";
  List<NamedEntityDescriptor> launcherDescriptors = new LinkedList<>();
  ByteBuffer bb = ByteBuffer.allocate(4);
  bb.putInt(0, 3);
  UserPayload customPayload = UserPayload.create(bb);
  launcherDescriptors.add(
      new NamedEntityDescriptor(customLauncherName, FakeContainerLauncher.class.getName())
          .setUserPayload(customPayload));
  launcherDescriptors
      .add(new NamedEntityDescriptor(TezConstants.getTezYarnServicePluginName(), null)
          .setUserPayload(userPayload));

  ContainerLaucherRouterForMultipleLauncherTest clr =
      new ContainerLaucherRouterForMultipleLauncherTest(appContext, tal, null,
          launcherDescriptors,
          true);
  try {
    clr.init(conf);
    clr.start();

    assertEquals(2, clr.getNumContainerLaunchers());
    assertTrue(clr.getYarnContainerLauncherCreated());
    assertFalse(clr.getUberContainerLauncherCreated());
    assertEquals(customLauncherName, clr.getContainerLauncherName(0));
    assertEquals(bb, clr.getContainerLauncherContext(0).getInitialUserPayload().getPayload());

    assertEquals(TezConstants.getTezYarnServicePluginName(), clr.getContainerLauncherName(1));
    Configuration confParsed = TezUtils
        .createConfFromUserPayload(clr.getContainerLauncherContext(1).getInitialUserPayload());
    assertEquals("testvalue", confParsed.get("testkey"));
  } finally {
    clr.stop();
  }
}
 

@Test(timeout = 5000)
public void testTaskSchedulerRouting() throws Exception {
  Configuration conf = new Configuration(false);
  UserPayload defaultPayload = TezUtils.createUserPayloadFromConf(conf);

  String customSchedulerName = "fakeScheduler";
  List<NamedEntityDescriptor> taskSchedulers = new LinkedList<>();
  ByteBuffer bb = ByteBuffer.allocate(4);
  bb.putInt(0, 3);
  UserPayload userPayload = UserPayload.create(bb);
  taskSchedulers.add(
      new NamedEntityDescriptor(customSchedulerName, FakeTaskScheduler.class.getName())
          .setUserPayload(userPayload));
  taskSchedulers.add(new NamedEntityDescriptor(TezConstants.getTezYarnServicePluginName(), null)
      .setUserPayload(defaultPayload));

  TSEHForMultipleSchedulersTest tseh =
      new TSEHForMultipleSchedulersTest(mockAppContext, mockClientService, mockEventHandler,
          mockSigMatcher, mockWebUIService, taskSchedulers, false);

  tseh.init(conf);
  tseh.start();

  // Verify that the YARN task scheduler is installed by default
  assertTrue(tseh.getYarnSchedulerCreated());
  assertFalse(tseh.getUberSchedulerCreated());
  assertEquals(2, tseh.getNumCreateInvocations());

  // Verify the order of the schedulers
  assertEquals(customSchedulerName, tseh.getTaskSchedulerName(0));
  assertEquals(TezConstants.getTezYarnServicePluginName(), tseh.getTaskSchedulerName(1));

  verify(tseh.getTestTaskScheduler(0)).initialize();
  verify(tseh.getTestTaskScheduler(0)).start();

  ApplicationId appId = ApplicationId.newInstance(1000, 1);
  TezDAGID dagId = TezDAGID.getInstance(appId, 1);
  TezVertexID vertexID = TezVertexID.getInstance(dagId, 1);
  TezTaskID taskId1 = TezTaskID.getInstance(vertexID, 1);
  TezTaskAttemptID attemptId11 = TezTaskAttemptID.getInstance(taskId1, 1);
  TezTaskID taskId2 = TezTaskID.getInstance(vertexID, 2);
  TezTaskAttemptID attemptId21 = TezTaskAttemptID.getInstance(taskId2, 1);

  Resource resource = Resource.newInstance(1024, 1);

  TaskAttempt mockTaskAttempt1 = mock(TaskAttempt.class);
  TaskAttempt mockTaskAttempt2 = mock(TaskAttempt.class);

  AMSchedulerEventTALaunchRequest launchRequest1 =
      new AMSchedulerEventTALaunchRequest(attemptId11, resource, mock(TaskSpec.class),
          mockTaskAttempt1, mock(TaskLocationHint.class), 1, mock(ContainerContext.class), 0, 0,
          0);

  tseh.handle(launchRequest1);

  verify(tseh.getTestTaskScheduler(0)).allocateTask(eq(mockTaskAttempt1), eq(resource),
      any(String[].class), any(String[].class), any(Priority.class), any(Object.class),
      eq(launchRequest1));

  AMSchedulerEventTALaunchRequest launchRequest2 =
      new AMSchedulerEventTALaunchRequest(attemptId21, resource, mock(TaskSpec.class),
          mockTaskAttempt2, mock(TaskLocationHint.class), 1, mock(ContainerContext.class), 1, 0,
          0);
  tseh.handle(launchRequest2);
  verify(tseh.getTestTaskScheduler(1)).allocateTask(eq(mockTaskAttempt2), eq(resource),
      any(String[].class), any(String[].class), any(Priority.class), any(Object.class),
      eq(launchRequest2));
}
 

@Test(timeout = 5000)
public void testCustomTaskSchedulerSetup() throws IOException {
  Configuration conf = new Configuration(false);
  conf.set("testkey", "testval");
  UserPayload defaultPayload = TezUtils.createUserPayloadFromConf(conf);

  String customSchedulerName = "fakeScheduler";
  List<NamedEntityDescriptor> taskSchedulers = new LinkedList<>();
  ByteBuffer bb = ByteBuffer.allocate(4);
  bb.putInt(0, 3);
  UserPayload userPayload = UserPayload.create(bb);
  taskSchedulers.add(
      new NamedEntityDescriptor(customSchedulerName, FakeTaskScheduler.class.getName())
          .setUserPayload(userPayload));
  taskSchedulers.add(new NamedEntityDescriptor(TezConstants.getTezYarnServicePluginName(), null)
      .setUserPayload(defaultPayload));

  TSEHForMultipleSchedulersTest tseh =
      new TSEHForMultipleSchedulersTest(mockAppContext, mockClientService, mockEventHandler,
          mockSigMatcher, mockWebUIService, taskSchedulers, false);

  tseh.init(conf);
  tseh.start();

  // Verify that the YARN task scheduler is installed by default
  assertTrue(tseh.getYarnSchedulerCreated());
  assertFalse(tseh.getUberSchedulerCreated());
  assertEquals(2, tseh.getNumCreateInvocations());

  // Verify the order of the schedulers
  assertEquals(customSchedulerName, tseh.getTaskSchedulerName(0));
  assertEquals(TezConstants.getTezYarnServicePluginName(), tseh.getTaskSchedulerName(1));

  // Verify the payload setup for the custom task scheduler
  assertNotNull(tseh.getTaskSchedulerContext(0));
  assertEquals(bb, tseh.getTaskSchedulerContext(0).getInitialUserPayload().getPayload());

  // Verify the payload on the yarn scheduler
  assertNotNull(tseh.getTaskSchedulerContext(1));
  Configuration parsed = TezUtils.createConfFromUserPayload(tseh.getTaskSchedulerContext(1).getInitialUserPayload());
  assertEquals("testval", parsed.get("testkey"));
}
 

@BeforeClass
public static void setup() throws Exception {

  extServiceTestHelper = new ExternalTezServiceTestHelper(TEST_ROOT_DIR);
  UserPayload userPayload =
      TezUtils.createUserPayloadFromConf(extServiceTestHelper.getConfForJobs());
  UserPayload userPayloadThrowError =
      ErrorPluginConfiguration.toUserPayload(ErrorPluginConfiguration.createThrowErrorConf());

  UserPayload userPayloadReportFatalErrorLauncher = ErrorPluginConfiguration
      .toUserPayload(ErrorPluginConfiguration.createReportFatalErrorConf(SUFFIX_LAUNCHER));
  UserPayload userPayloadReportFatalErrorTaskComm = ErrorPluginConfiguration
      .toUserPayload(ErrorPluginConfiguration.createReportFatalErrorConf(SUFFIX_TASKCOMM));
  UserPayload userPayloadReportFatalErrorScheduler = ErrorPluginConfiguration
      .toUserPayload(ErrorPluginConfiguration.createReportFatalErrorConf(SUFFIX_SCHEDULER));

  UserPayload userPayloadReportNonFatalErrorLauncher = ErrorPluginConfiguration
      .toUserPayload(ErrorPluginConfiguration.createReportNonFatalErrorConf(SUFFIX_LAUNCHER));
  UserPayload userPayloadReportNonFatalErrorTaskComm = ErrorPluginConfiguration
      .toUserPayload(ErrorPluginConfiguration.createReportNonFatalErrorConf(SUFFIX_TASKCOMM));
  UserPayload userPayloadReportNonFatalErrorScheduler = ErrorPluginConfiguration
      .toUserPayload(ErrorPluginConfiguration.createReportNonFatalErrorConf(SUFFIX_SCHEDULER));

  TaskSchedulerDescriptor[] taskSchedulerDescriptors = new TaskSchedulerDescriptor[]{
      TaskSchedulerDescriptor
          .create(EXT_PUSH_ENTITY_NAME, TezTestServiceTaskSchedulerService.class.getName())
          .setUserPayload(userPayload),
      TaskSchedulerDescriptor.create(EXT_THROW_ERROR_ENTITY_NAME,
          TezTestServiceTaskSchedulerServiceWithErrors.class.getName()).setUserPayload(
          userPayloadThrowError),
      TaskSchedulerDescriptor.create(EXT_REPORT_FATAL_ERROR_ENTITY_NAME,
          TezTestServiceTaskSchedulerServiceWithErrors.class.getName()).setUserPayload(
          userPayloadReportFatalErrorScheduler),
      TaskSchedulerDescriptor.create(EXT_REPORT_NON_FATAL_ERROR_ENTITY_NAME,
          TezTestServiceTaskSchedulerServiceWithErrors.class.getName()).setUserPayload(
          userPayloadReportNonFatalErrorScheduler),
  };

  ContainerLauncherDescriptor[] containerLauncherDescriptors = new ContainerLauncherDescriptor[]{
      ContainerLauncherDescriptor
          .create(EXT_PUSH_ENTITY_NAME, TezTestServiceNoOpContainerLauncher.class.getName())
          .setUserPayload(userPayload),
      ContainerLauncherDescriptor.create(EXT_THROW_ERROR_ENTITY_NAME,
          TezTestServiceContainerLauncherWithErrors.class.getName()).setUserPayload(userPayloadThrowError),
      ContainerLauncherDescriptor.create(EXT_REPORT_FATAL_ERROR_ENTITY_NAME,
          TezTestServiceContainerLauncherWithErrors.class.getName()).setUserPayload(userPayloadReportFatalErrorLauncher),
      ContainerLauncherDescriptor.create(EXT_REPORT_NON_FATAL_ERROR_ENTITY_NAME,
          TezTestServiceContainerLauncherWithErrors.class.getName()).setUserPayload(userPayloadReportNonFatalErrorLauncher)
  };

  TaskCommunicatorDescriptor[] taskCommunicatorDescriptors = new TaskCommunicatorDescriptor[]{
      TaskCommunicatorDescriptor
          .create(EXT_PUSH_ENTITY_NAME, TezTestServiceTaskCommunicatorImpl.class.getName())
          .setUserPayload(userPayload),
      TaskCommunicatorDescriptor.create(EXT_THROW_ERROR_ENTITY_NAME,
          TezTestServiceTaskCommunicatorWithErrors.class.getName()).setUserPayload(userPayloadThrowError),
      TaskCommunicatorDescriptor.create(EXT_REPORT_FATAL_ERROR_ENTITY_NAME,
          TezTestServiceTaskCommunicatorWithErrors.class.getName()).setUserPayload(userPayloadReportFatalErrorTaskComm),
      TaskCommunicatorDescriptor.create(EXT_REPORT_NON_FATAL_ERROR_ENTITY_NAME,
          TezTestServiceTaskCommunicatorWithErrors.class.getName()).setUserPayload(userPayloadReportNonFatalErrorTaskComm)
  };

  servicePluginsDescriptor = ServicePluginsDescriptor.create(true, true,
      taskSchedulerDescriptors, containerLauncherDescriptors, taskCommunicatorDescriptors);

  extServiceTestHelper.setupSharedTezClient(servicePluginsDescriptor);

  // Generate the join data set used for each run.
  // Can a timeout be enforced here ?
  Path dataPath1 = new Path(SRC_DATA_DIR, "inPath1");
  Path dataPath2 = new Path(SRC_DATA_DIR, "inPath2");
  extServiceTestHelper
      .setupHashJoinData(SRC_DATA_DIR, dataPath1, dataPath2, HASH_JOIN_EXPECTED_RESULT_PATH, HASH_JOIN_OUTPUT_PATH);

  extServiceTestHelper.shutdownSharedTezClient();
}
 

@Test(timeout = 5000)
public void testMultipleTaskComms() throws IOException, TezException {

  AppContext appContext = mock(AppContext.class);
  TaskHeartbeatHandler thh = mock(TaskHeartbeatHandler.class);
  ContainerHeartbeatHandler chh = mock(ContainerHeartbeatHandler.class);
  Configuration conf = new Configuration(false);
  conf.set("testkey", "testvalue");
  UserPayload defaultPayload = TezUtils.createUserPayloadFromConf(conf);

  String customTaskCommName = "customTaskComm";
  List<NamedEntityDescriptor> taskCommDescriptors = new LinkedList<>();
  ByteBuffer bb = ByteBuffer.allocate(4);
  bb.putInt(0, 3);
  UserPayload customPayload = UserPayload.create(bb);
  taskCommDescriptors.add(
      new NamedEntityDescriptor(customTaskCommName, FakeTaskComm.class.getName())
          .setUserPayload(customPayload));
  taskCommDescriptors
      .add(new NamedEntityDescriptor(TezConstants.getTezYarnServicePluginName(), null).setUserPayload(defaultPayload));

  TaskCommManagerForMultipleCommTest tcm =
      new TaskCommManagerForMultipleCommTest(appContext, thh, chh, taskCommDescriptors);

  try {
    tcm.init(new Configuration(false));
    tcm.start();

    assertEquals(2, tcm.getNumTaskComms());
    assertTrue(tcm.getYarnTaskCommCreated());
    assertFalse(tcm.getUberTaskCommCreated());

    assertEquals(customTaskCommName, tcm.getTaskCommName(0));
    assertEquals(bb, tcm.getTaskCommContext(0).getInitialUserPayload().getPayload());

    assertEquals(TezConstants.getTezYarnServicePluginName(), tcm.getTaskCommName(1));
    Configuration confParsed = TezUtils
        .createConfFromUserPayload(tcm.getTaskCommContext(1).getInitialUserPayload());
    assertEquals("testvalue", confParsed.get("testkey"));
  } finally {
    tcm.stop();
  }
}