Java Code Examples for org.apache.tez.dag.api.UserPayload#create()

@Override
public void initialize(CartesianProductConfigProto config) throws Exception {
  this.numPartitions = Ints.toArray(config.getNumPartitionsList());
  this.sources = config.getSourcesList();
  this.positionId = sources.indexOf(getContext().getSourceVertexName());

  if (config.hasFilterClassName()) {
    UserPayload userPayload = config.hasFilterUserPayload()
      ? UserPayload.create(ByteBuffer.wrap(config.getFilterUserPayload().toByteArray())) : null;
    try {
      filter = ReflectionUtils.createClazzInstance(config.getFilterClassName(),
        new Class[]{UserPayload.class}, new UserPayload[]{userPayload});
    } catch (TezReflectionException e) {
      throw e;
    }
  }
  generateTaskIdMapping();
}
 

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 UserPayload toUserPayload(ErrorPluginConfiguration conf) throws IOException {
  ByteArrayOutputStream baos = new ByteArrayOutputStream();
  ObjectOutputStream oos = new ObjectOutputStream(baos);
  oos.writeObject(conf.kv);
  oos.close();
  UserPayload userPayload = UserPayload.create(ByteBuffer.wrap(baos.toByteArray()));
  return userPayload;
}
 

/**
 * create a DAG with 2 vertices (v1 --> v2), set payload on Input/Output/Processor/VertexManagerPlugin to
 * control where throw exception
 * 
 * @param exLocation
 * @return
 * @throws IOException
 */
private DAG createDAG(ExceptionLocation exLocation) throws IOException {
  DAG dag = DAG.create("dag_" + exLocation.name());
  UserPayload payload =
      UserPayload.create(ByteBuffer.wrap(exLocation.name().getBytes()));
  Vertex v1 =
      Vertex.create("v1", ProcessorWithException.getProcDesc(payload), 1);
  InputDescriptor inputDesc = InputWithException.getInputDesc(payload);
  InputInitializerDescriptor iiDesc =
      InputInitializerWithException.getIIDesc(payload);
  v1.addDataSource("input",
      DataSourceDescriptor.create(inputDesc, iiDesc, null));
  v1.setVertexManagerPlugin(RootInputVertexManagerWithException
      .getVMDesc(exLocation));

  Vertex v2 = 
      Vertex.create("v2", DoNothingProcessor.getProcDesc(), 1);
  v2.addDataSource("input2",
      DataSourceDescriptor.create(InputDescriptor.create(NoOpInput.class.getName()),
        InputInitializerWithException2.getIIDesc(payload), null));

  dag.addVertex(v1)
    .addVertex(v2);
  if (exLocation.name().startsWith("EM_")) {
    dag.addEdge(Edge.create(v1, v2, EdgeProperty.create(
        EdgeManagerPluginDescriptor.create(CustomEdgeManager.class.getName())
          .setUserPayload(payload),
        DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
        OutputWithException.getOutputDesc(payload), InputWithException.getInputDesc(payload))));
  } else {
    // set Customized VertexManager here, it can't been used for CustomEdge
    v2.setVertexManagerPlugin(InputReadyVertexManagerWithException.getVMDesc(exLocation));
    dag.addEdge(Edge.create(v1, v2, EdgeProperty.create(DataMovementType.ONE_TO_ONE,
        DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
        OutputWithException.getOutputDesc(payload), InputWithException.getInputDesc(payload))));
  }

  return dag;
}
 

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;
}
 

public static DAG createDAG(String name,
    Configuration conf) throws Exception {
  UserPayload payload = UserPayload.create(null);
  int taskCount = MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS_DEFAULT;
  if (conf != null) {
    taskCount = conf.getInt(MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS, MULTI_ATTEMPT_DAG_VERTEX_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);

  // Make each vertex manager fail on appropriate attempt
  v1.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
      FailOnAttemptVertexManagerPlugin.class.getName())
      .setUserPayload(UserPayload.create(ByteBuffer.wrap(new String("1").getBytes()))));
  v2.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
      FailOnAttemptVertexManagerPlugin.class.getName())
      .setUserPayload(UserPayload.create(ByteBuffer.wrap(new String("2").getBytes()))));
  v3.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
      FailOnAttemptVertexManagerPlugin.class.getName())
      .setUserPayload(UserPayload.create(ByteBuffer.wrap(new String("3").getBytes()))));
  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(v2, v3,
      EdgeProperty.create(DataMovementType.SCATTER_GATHER,
          DataSourceType.PERSISTED,
          SchedulingType.SEQUENTIAL,
          TestOutput.getOutputDesc(payload),
          TestInput.getInputDesc(payload))));
  return dag;
}
 

/**
 *  v1  v2
 *   \  /
 *    v3
 *   /  \
 *  v4  v5
 *   \  /
 *    v6
 * @param name
 * @param conf
 * @return
 * @throws Exception
 */
public static DAG createDAGForVertexOrder(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);
  Vertex v3 = Vertex.create("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v4 = Vertex.create("v4", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v5 = Vertex.create("v5", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
  Vertex v6 = Vertex.create("v6", TestProcessor.getProcDesc(payload), taskCount, defaultResource);

  // add vertex not in the topological order, since we are using this dag for testing vertex topological order
  dag.addVertex(v4)
    .addVertex(v5)
    .addVertex(v6)
    .addVertex(v1)
    .addVertex(v2)
    .addVertex(v3)
    .addEdge(Edge.create(v1, v3,
        EdgeProperty.create(DataMovementType.SCATTER_GATHER,
            DataSourceType.PERSISTED,
            SchedulingType.SEQUENTIAL,
            TestOutput.getOutputDesc(payload),
            TestInput.getInputDesc(payload))))
    .addEdge(Edge.create(v2, v3,
        EdgeProperty.create(DataMovementType.SCATTER_GATHER,
            DataSourceType.PERSISTED,
            SchedulingType.SEQUENTIAL,
            TestOutput.getOutputDesc(payload),
            TestInput.getInputDesc(payload))))
    .addEdge(Edge.create(v3, v4,
        EdgeProperty.create(DataMovementType.SCATTER_GATHER,
            DataSourceType.PERSISTED,
            SchedulingType.SEQUENTIAL,
            TestOutput.getOutputDesc(payload),
            TestInput.getInputDesc(payload))))
    .addEdge(Edge.create(v3, v5,
        EdgeProperty.create(DataMovementType.SCATTER_GATHER,
            DataSourceType.PERSISTED,
            SchedulingType.SEQUENTIAL,
            TestOutput.getOutputDesc(payload),
            TestInput.getInputDesc(payload))))
    .addEdge(Edge.create(v4, v6,
        EdgeProperty.create(DataMovementType.SCATTER_GATHER,
            DataSourceType.PERSISTED,
            SchedulingType.SEQUENTIAL,
            TestOutput.getOutputDesc(payload),
            TestInput.getInputDesc(payload))))
    .addEdge(Edge.create(v5, v6,
        EdgeProperty.create(DataMovementType.SCATTER_GATHER,
            DataSourceType.PERSISTED,
            SchedulingType.SEQUENTIAL,
            TestOutput.getOutputDesc(payload),
            TestInput.getInputDesc(payload))));

  return dag;
}
 

private DAG createDAG(FileSystem fs, TezConfiguration tezConf,
    Path stagingDir, boolean doLocalityCheck) throws IOException, YarnException {

  int numBroadcastTasks = 2;
  int numOneToOneTasks = 3;
  if (doLocalityCheck) {
    YarnClient yarnClient = YarnClient.createYarnClient();
    yarnClient.init(tezConf);
    yarnClient.start();
    int numNMs = yarnClient.getNodeReports(NodeState.RUNNING).size();
    yarnClient.stop();
    // create enough 1-1 tasks to run in parallel
    numOneToOneTasks = numNMs - numBroadcastTasks - 1;// 1 AM
    if (numOneToOneTasks < 1) {
      numOneToOneTasks = 1;
    }
  }
  byte[] procByte = {(byte) (doLocalityCheck ? 1 : 0), 1};
  UserPayload procPayload = UserPayload.create(ByteBuffer.wrap(procByte));

  System.out.println("Using " + numOneToOneTasks + " 1-1 tasks");

  Vertex broadcastVertex = Vertex.create("Broadcast", ProcessorDescriptor.create(
      InputProcessor.class.getName()), numBroadcastTasks);
  
  Vertex inputVertex = Vertex.create("Input", ProcessorDescriptor.create(
      InputProcessor.class.getName()).setUserPayload(procPayload), numOneToOneTasks);

  Vertex oneToOneVertex = Vertex.create("OneToOne",
      ProcessorDescriptor.create(
          OneToOneProcessor.class.getName()).setUserPayload(procPayload));
  oneToOneVertex.setVertexManagerPlugin(
      VertexManagerPluginDescriptor.create(InputReadyVertexManager.class.getName()));

  UnorderedKVEdgeConfig edgeConf = UnorderedKVEdgeConfig
      .newBuilder(Text.class.getName(), IntWritable.class.getName())
      .setFromConfiguration(tezConf).build();

  DAG dag = DAG.create("BroadcastAndOneToOneExample");
  dag.addVertex(inputVertex)
      .addVertex(broadcastVertex)
      .addVertex(oneToOneVertex)
      .addEdge(
          Edge.create(inputVertex, oneToOneVertex, edgeConf.createDefaultOneToOneEdgeProperty()))
      .addEdge(
          Edge.create(broadcastVertex, oneToOneVertex,
              edgeConf.createDefaultBroadcastEdgeProperty()));
  return dag;
}
 

@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();
  }
}
 

@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"));
}
 

@Test(timeout = 5000)
public void testSerializedPayload() throws IOException {

  Configuration conf = new Configuration(false);
  conf.setBoolean(MRJobConfig.MR_TEZ_INPUT_INITIALIZER_SERIALIZE_EVENT_PAYLOAD, true);
  ByteString confByteString = TezUtils.createByteStringFromConf(conf);
  InputSplit split1 = new InputSplitForTest(1);
  InputSplit split2 = new InputSplitForTest(2);
  MRSplitProto proto1 = MRInputHelpers.createSplitProto(split1);
  MRSplitProto proto2 = MRInputHelpers.createSplitProto(split2);
  MRSplitsProto.Builder splitsProtoBuilder = MRSplitsProto.newBuilder();
  splitsProtoBuilder.addSplits(proto1);
  splitsProtoBuilder.addSplits(proto2);
  MRInputUserPayloadProto.Builder payloadProto = MRInputUserPayloadProto.newBuilder();
  payloadProto.setSplits(splitsProtoBuilder.build());
  payloadProto.setConfigurationBytes(confByteString);
  UserPayload userPayload =
      UserPayload.create(payloadProto.build().toByteString().asReadOnlyByteBuffer());

  InputInitializerContext context = new TezTestUtils.TezRootInputInitializerContextForTest(userPayload,
      new Configuration(false));
  MRInputSplitDistributor splitDist = new MRInputSplitDistributor(context);

  List<Event> events = splitDist.initialize();

  assertEquals(3, events.size());
  assertTrue(events.get(0) instanceof InputUpdatePayloadEvent);
  assertTrue(events.get(1) instanceof InputDataInformationEvent);
  assertTrue(events.get(2) instanceof InputDataInformationEvent);

  InputDataInformationEvent diEvent1 = (InputDataInformationEvent) (events.get(1));
  InputDataInformationEvent diEvent2 = (InputDataInformationEvent) (events.get(2));

  assertNull(diEvent1.getDeserializedUserPayload());
  assertNull(diEvent2.getDeserializedUserPayload());

  assertNotNull(diEvent1.getUserPayload());
  assertNotNull(diEvent2.getUserPayload());

  MRSplitProto event1Proto = MRSplitProto.parseFrom(ByteString.copyFrom(diEvent1.getUserPayload()));
  InputSplit is1 = MRInputUtils.getOldSplitDetailsFromEvent(event1Proto, new Configuration());
  assertTrue(is1 instanceof InputSplitForTest);
  assertEquals(1, ((InputSplitForTest) is1).identifier);

  MRSplitProto event2Proto = MRSplitProto.parseFrom(ByteString.copyFrom(diEvent2.getUserPayload()));
  InputSplit is2 = MRInputUtils.getOldSplitDetailsFromEvent(event2Proto, new Configuration());
  assertTrue(is2 instanceof InputSplitForTest);
  assertEquals(2, ((InputSplitForTest) is2).identifier);
}
 

public UserPayload toUserPayload(TezConfiguration conf) throws IOException {
  return UserPayload.create(ByteBuffer.wrap(toProto(conf).toByteArray()));
}
 

@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 testEventRouting() throws Exception {
  Configuration conf = new Configuration(false);
  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(TezConstants.getTezYarnServicePluginName(), clr.getContainerLauncherName(1));

    verify(clr.getTestContainerLauncher(0)).initialize();
    verify(clr.getTestContainerLauncher(0)).start();
    verify(clr.getTestContainerLauncher(1)).initialize();
    verify(clr.getTestContainerLauncher(1)).start();

    ContainerLaunchContext clc1 = mock(ContainerLaunchContext.class);
    Container container1 = mock(Container.class);

    ContainerLaunchContext clc2 = mock(ContainerLaunchContext.class);
    Container container2 = mock(Container.class);

    ContainerLauncherLaunchRequestEvent launchRequestEvent1 =
        new ContainerLauncherLaunchRequestEvent(clc1, container1, 0, 0, 0);
    ContainerLauncherLaunchRequestEvent launchRequestEvent2 =
        new ContainerLauncherLaunchRequestEvent(clc2, container2, 1, 0, 0);

    clr.handle(launchRequestEvent1);


    ArgumentCaptor<ContainerLaunchRequest> captor =
        ArgumentCaptor.forClass(ContainerLaunchRequest.class);
    verify(clr.getTestContainerLauncher(0)).launchContainer(captor.capture());
    assertEquals(1, captor.getAllValues().size());
    ContainerLaunchRequest launchRequest1 = captor.getValue();
    assertEquals(clc1, launchRequest1.getContainerLaunchContext());

    clr.handle(launchRequestEvent2);
    captor = ArgumentCaptor.forClass(ContainerLaunchRequest.class);
    verify(clr.getTestContainerLauncher(1)).launchContainer(captor.capture());
    assertEquals(1, captor.getAllValues().size());
    ContainerLaunchRequest launchRequest2 = captor.getValue();
    assertEquals(clc2, launchRequest2.getContainerLaunchContext());

  } finally {
    clr.stop();
    verify(clr.getTestContainerLauncher(0)).shutdown();
    verify(clr.getTestContainerLauncher(1)).shutdown();
  }
}
 

public UserPayload toUserPayload() {
  return UserPayload.create(ByteBuffer.wrap(Integer.toString(timeToSleepMS).getBytes(
      charSet)));
}
 

UserPayload getBytePayload() {
  return UserPayload.create(ByteBuffer.wrap(
      getConfigPayload().toByteArray()));
}
 

@Test(timeout = 5000)
public void testDeserializedPayload() throws IOException {

  Configuration conf = new Configuration(false);
  conf.setBoolean(MRJobConfig.MR_TEZ_INPUT_INITIALIZER_SERIALIZE_EVENT_PAYLOAD, false);
  ByteString confByteString = TezUtils.createByteStringFromConf(conf);
  InputSplit split1 = new InputSplitForTest(1);
  InputSplit split2 = new InputSplitForTest(2);
  MRSplitProto proto1 = MRInputHelpers.createSplitProto(split1);
  MRSplitProto proto2 = MRInputHelpers.createSplitProto(split2);
  MRSplitsProto.Builder splitsProtoBuilder = MRSplitsProto.newBuilder();
  splitsProtoBuilder.addSplits(proto1);
  splitsProtoBuilder.addSplits(proto2);
  MRInputUserPayloadProto.Builder payloadProto = MRInputUserPayloadProto.newBuilder();
  payloadProto.setSplits(splitsProtoBuilder.build());
  payloadProto.setConfigurationBytes(confByteString);
  UserPayload userPayload =
      UserPayload.create(payloadProto.build().toByteString().asReadOnlyByteBuffer());

  InputInitializerContext context = new TezTestUtils.TezRootInputInitializerContextForTest(userPayload,
      new Configuration(false));
  MRInputSplitDistributor splitDist = new MRInputSplitDistributor(context);

  List<Event> events = splitDist.initialize();

  assertEquals(3, events.size());
  assertTrue(events.get(0) instanceof InputUpdatePayloadEvent);
  assertTrue(events.get(1) instanceof InputDataInformationEvent);
  assertTrue(events.get(2) instanceof InputDataInformationEvent);

  InputDataInformationEvent diEvent1 = (InputDataInformationEvent) (events.get(1));
  InputDataInformationEvent diEvent2 = (InputDataInformationEvent) (events.get(2));

  assertNull(diEvent1.getUserPayload());
  assertNull(diEvent2.getUserPayload());

  assertNotNull(diEvent1.getDeserializedUserPayload());
  assertNotNull(diEvent2.getDeserializedUserPayload());

  assertTrue(diEvent1.getDeserializedUserPayload() instanceof InputSplitForTest);
  assertEquals(1, ((InputSplitForTest) diEvent1.getDeserializedUserPayload()).identifier);

  assertTrue(diEvent2.getDeserializedUserPayload() instanceof InputSplitForTest);
  assertEquals(2, ((InputSplitForTest) diEvent2.getDeserializedUserPayload()).identifier);
}
 

@Test(timeout = 5000)
public void testEventRouting() throws Exception {

  AppContext appContext = mock(AppContext.class, RETURNS_DEEP_STUBS);
  NodeId nodeId = NodeId.newInstance("host1", 3131);
  when(appContext.getAllContainers().get(any(ContainerId.class)).getContainer().getNodeId())
      .thenReturn(nodeId);
  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());

    verify(tcm.getTestTaskComm(0)).initialize();
    verify(tcm.getTestTaskComm(0)).start();
    verify(tcm.getTestTaskComm(1)).initialize();
    verify(tcm.getTestTaskComm(1)).start();


    ContainerId containerId1 = mock(ContainerId.class);
    tcm.registerRunningContainer(containerId1, 0);
    verify(tcm.getTestTaskComm(0)).registerRunningContainer(eq(containerId1), eq("host1"),
        eq(3131));

    ContainerId containerId2 = mock(ContainerId.class);
    tcm.registerRunningContainer(containerId2, 1);
    verify(tcm.getTestTaskComm(1)).registerRunningContainer(eq(containerId2), eq("host1"),
        eq(3131));

  } finally {
    tcm.stop();
    verify(tcm.getTaskCommunicator(0).getTaskCommunicator()).shutdown();
    verify(tcm.getTaskCommunicator(1).getTaskCommunicator()).shutdown();
  }
}
 

public TezRootInputInitializerContextForTest(UserPayload payload, Configuration vertexConfig) throws IOException {
  appId = ApplicationId.newInstance(1000, 200);
  this.payload = payload == null ? UserPayload.create(null) : payload;
  this.vertexConfig = vertexConfig;
}
 

/**
 * Convert a Configuration to a {@link org.apache.tez.dag.api.UserPayload} </p>
 *
 *
 * @param conf configuration to be converted
 * @return an instance of {@link org.apache.tez.dag.api.UserPayload}
 * @throws java.io.IOException
 */
public static UserPayload createUserPayloadFromConf(Configuration conf) throws IOException {
  return UserPayload.create(ByteBuffer.wrap(createByteStringFromConf(conf).toByteArray()));
}