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

@Test(timeout = 5000)
public void testInvalidPhysicalOutputCount() throws Exception {
  EventHandler mockEventHandler = mock(EventHandler.class);
  Edge edge = new Edge(EdgeProperty.create(
      EdgeManagerPluginDescriptor.create(CustomEdgeManagerWithInvalidReturnValue.class.getName())
        .setUserPayload(new CustomEdgeManagerWithInvalidReturnValue.EdgeManagerConfig(1,-1,1,1).toUserPayload()),
      DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(""),
      InputDescriptor.create("")), mockEventHandler, new TezConfiguration());
  TezVertexID v1Id = createVertexID(1);
  TezVertexID v2Id = createVertexID(2);
  edge.setSourceVertex(mockVertex("v1", v1Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.setDestinationVertex(mockVertex("v2", v2Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.initialize();
  try {
    edge.getSourceSpec(0);
    Assert.fail();
  } catch (AMUserCodeException e) {
    e.printStackTrace();
    assertTrue(e.getCause().getMessage().contains("PhysicalOutputCount should not be negative"));
  }
}
 

@Test(timeout = 5000)
public void testInvalidPhysicalInputCount() throws Exception {
  EventHandler mockEventHandler = mock(EventHandler.class);
  Edge edge = new Edge(EdgeProperty.create(
      EdgeManagerPluginDescriptor.create(CustomEdgeManagerWithInvalidReturnValue.class.getName())
        .setUserPayload(new CustomEdgeManagerWithInvalidReturnValue.EdgeManagerConfig(-1,1,1,1).toUserPayload()),
      DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(""),
      InputDescriptor.create("")), mockEventHandler, new TezConfiguration());
  TezVertexID v1Id = createVertexID(1);
  TezVertexID v2Id = createVertexID(2);
  edge.setSourceVertex(mockVertex("v1", v1Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.setDestinationVertex(mockVertex("v2", v2Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.initialize();
  try {
    edge.getDestinationSpec(0);
    Assert.fail();
  } catch (AMUserCodeException e) {
    e.printStackTrace();
    assertTrue(e.getCause().getMessage().contains("PhysicalInputCount should not be negative"));
  }
}
 

@Test(timeout = 5000)
public void testInvalidSourceTaskIndex() throws Exception {
  EventHandler mockEventHandler = mock(EventHandler.class);
  Edge edge = new Edge(EdgeProperty.create(
      EdgeManagerPluginDescriptor.create(CustomEdgeManagerWithInvalidReturnValue.class.getName())
        .setUserPayload(new CustomEdgeManagerWithInvalidReturnValue.EdgeManagerConfig(1,1,1,-1).toUserPayload()),
      DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(""),
      InputDescriptor.create("")), mockEventHandler, new TezConfiguration());
  TezVertexID v1Id = createVertexID(1);
  TezVertexID v2Id = createVertexID(2);
  edge.setSourceVertex(mockVertex("v1", v1Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.setDestinationVertex(mockVertex("v2", v2Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.initialize();
  try {
    TezEvent ireEvent = new TezEvent(InputReadErrorEvent.create("diag", 0, 1),
        new EventMetaData(EventProducerConsumerType.INPUT, "v2", "v1",
            TezTaskAttemptID.getInstance(TezTaskID.getInstance(v2Id, 1), 1)));
    edge.sendTezEventToSourceTasks(ireEvent);
    Assert.fail();
  } catch (AMUserCodeException e) {
    e.printStackTrace();
    assertTrue(e.getCause().getMessage().contains("SourceTaskIndex should not be negative"));
  }
}
 

@Test(timeout = 5000)
public void testInvalidConsumerNumber() throws Exception {
  EventHandler mockEventHandler = mock(EventHandler.class);
  Edge edge = new Edge(EdgeProperty.create(
      EdgeManagerPluginDescriptor.create(CustomEdgeManagerWithInvalidReturnValue.class.getName())
        .setUserPayload(new CustomEdgeManagerWithInvalidReturnValue.EdgeManagerConfig(1,1,0,1).toUserPayload()),
      DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(""),
      InputDescriptor.create("")), mockEventHandler, new TezConfiguration());
  TezVertexID v1Id = createVertexID(1);
  TezVertexID v2Id = createVertexID(2);
  edge.setSourceVertex(mockVertex("v1", v1Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.setDestinationVertex(mockVertex("v2", v2Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.initialize();
  try {
    TezEvent ireEvent = new TezEvent(InputReadErrorEvent.create("diag", 0, 1),
        new EventMetaData(EventProducerConsumerType.INPUT, "v2", "v1",
            TezTaskAttemptID.getInstance(TezTaskID.getInstance(v2Id, 1), 1)));
    edge.sendTezEventToSourceTasks(ireEvent);
    Assert.fail();
  } catch (AMUserCodeException e) {
    e.printStackTrace();
    assertTrue(e.getCause().getMessage().contains("ConsumerTaskNum must be positive"));
  }
}
 

/**
 * This is a convenience method for the typical usage of this edge, and creates an instance of
 * {@link org.apache.tez.dag.api.EdgeProperty} which is likely to be used. </p>
 * If custom edge properties are required, the methods to get the relevant payloads should be
 * used. </p>
 * * In this case - DataMovementType.BROADCAST, EdgeProperty.DataSourceType.PERSISTED,
 * EdgeProperty.SchedulingType.SEQUENTIAL
 *
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultBroadcastEdgeProperty() {
  EdgeProperty edgeProperty = EdgeProperty.create(EdgeProperty.DataMovementType.BROADCAST,
      EdgeProperty.DataSourceType.PERSISTED, EdgeProperty.SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(
          getOutputClassName()).setUserPayload(getOutputPayload()),
      InputDescriptor.create(
          getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

/**
 * This is a convenience method for the typical usage of this edge, and creates an instance of
 * {@link org.apache.tez.dag.api.EdgeProperty} which is likely to be used. </p>
 * If custom edge properties are required, the methods to get the relevant payloads should be
 * used. </p>
 * * In this case - DataMovementType.ONE_TO_ONE, EdgeProperty.DataSourceType.PERSISTED,
 * EdgeProperty.SchedulingType.SEQUENTIAL
 *
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultOneToOneEdgeProperty() {
  EdgeProperty edgeProperty = EdgeProperty.create(EdgeProperty.DataMovementType.ONE_TO_ONE,
      EdgeProperty.DataSourceType.PERSISTED, EdgeProperty.SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(
          getOutputClassName()).setUserPayload(getOutputPayload()),
      InputDescriptor.create(
          getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

/**
 * This is a convenience method for creating an Edge descriptor based on the specified
 * EdgeManagerDescriptor.
 *
 * @param edgeManagerDescriptor the custom edge specification
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultCustomEdgeProperty(EdgeManagerPluginDescriptor edgeManagerDescriptor) {
  Objects.requireNonNull(edgeManagerDescriptor, "EdgeManagerDescriptor cannot be null");
  EdgeProperty edgeProperty =
      EdgeProperty.create(edgeManagerDescriptor, EdgeProperty.DataSourceType.PERSISTED,
          EdgeProperty.SchedulingType.SEQUENTIAL,
          OutputDescriptor.create(getOutputClassName()).setUserPayload(getOutputPayload()),
          InputDescriptor.create(getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

/**
 * This is a convenience method for the typical usage of this edge, and creates an instance of
 * {@link org.apache.tez.dag.api.EdgeProperty} which is likely to be used. </p>
 * * In this case - DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED,
 * EdgeProperty.SchedulingType.SEQUENTIAL
 *
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultEdgeProperty() {
  EdgeProperty edgeProperty = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER,
      EdgeProperty.DataSourceType.PERSISTED, EdgeProperty.SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(
          getOutputClassName()).setUserPayload(getOutputPayload()),
      InputDescriptor.create(
          getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

/**
 * This is a convenience method for creating an Edge descriptor based on the specified
 * EdgeManagerDescriptor.
 *
 * @param edgeManagerDescriptor the custom edge specification
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultCustomEdgeProperty(EdgeManagerPluginDescriptor edgeManagerDescriptor) {
  Objects.requireNonNull(edgeManagerDescriptor, "EdgeManagerDescriptor cannot be null");
  EdgeProperty edgeProperty =
      EdgeProperty.create(edgeManagerDescriptor, EdgeProperty.DataSourceType.PERSISTED,
          EdgeProperty.SchedulingType.SEQUENTIAL,
          OutputDescriptor.create(getOutputClassName()).setUserPayload(getOutputPayload()),
          InputDescriptor.create(getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

/**
 * This is a convenience method for creating an Edge descriptor based on the specified
 * EdgeManagerDescriptor.
 *
 * @param edgeManagerDescriptor the custom edge specification
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultCustomEdgeProperty(EdgeManagerPluginDescriptor edgeManagerDescriptor) {
  Objects.requireNonNull(edgeManagerDescriptor, "EdgeManagerDescriptor cannot be null");
  EdgeProperty edgeProperty =
      EdgeProperty.create(edgeManagerDescriptor, EdgeProperty.DataSourceType.PERSISTED,
          EdgeProperty.SchedulingType.SEQUENTIAL,
          OutputDescriptor.create(getOutputClassName()).setUserPayload(getOutputPayload()),
          InputDescriptor.create(getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

@Test
public void testZeroTasksSendsConfigured() throws IOException {
  Configuration conf = new Configuration();
  ShuffleVertexManagerBase manager = null;

  HashMap<String, EdgeProperty> mockInputVertices = new HashMap<String, EdgeProperty>();
  String r1 = "R1";
  EdgeProperty eProp1 = EdgeProperty.create(
      EdgeProperty.DataMovementType.SCATTER_GATHER,
      EdgeProperty.DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));

  final String mockManagedVertexId = "R2";
  mockInputVertices.put(r1, eProp1);

  final VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
  when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
  when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
  when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(0);

  // check initialization
  manager = createManager(conf, mockContext, 0.001f, 0.001f);

  final List<Integer> scheduledTasks = Lists.newLinkedList();
  doAnswer(new ScheduledTasksAnswer(scheduledTasks)).when(
      mockContext).scheduleTasks(anyList());

  manager.onVertexStarted(emptyCompletions);
  manager.onVertexStateUpdated(new VertexStateUpdate(r1, VertexState.CONFIGURED));
  Assert.assertEquals(1, manager.bipartiteSources);
  Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
  Assert.assertEquals(0, manager.totalNumBipartiteSourceTasks);
  Assert.assertEquals(0, manager.pendingTasks.size()); // no tasks scheduled
  Assert.assertEquals(0, scheduledTasks.size());
  verify(mockContext).doneReconfiguringVertex();
}
 

@Test (timeout=5000)
public void testBasicScatterGather() throws Exception {
  HashMap<String, EdgeProperty> mockInputVertices = 
      new HashMap<String, EdgeProperty>();
  String mockSrcVertexId1 = "Vertex1";
  EdgeProperty eProp1 = EdgeProperty.create(
      EdgeProperty.DataMovementType.SCATTER_GATHER,
      EdgeProperty.DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));
  
  String mockManagedVertexId = "Vertex";
  
  VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
  when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
  when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
  when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(2);
  when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(3);
  mockInputVertices.put(mockSrcVertexId1, eProp1);

  InputReadyVertexManager manager = new InputReadyVertexManager(mockContext);
  manager.initialize();
  verify(mockContext, times(1)).vertexReconfigurationPlanned();
  // source vertex configured
  manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
  verify(mockContext, times(1)).doneReconfiguringVertex();
  verify(mockContext, times(0)).scheduleTasks(requestCaptor.capture());
  // then own vertex started
  manager.onVertexStarted(Collections.singletonList(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 0)));
  manager.onSourceTaskCompleted(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 1));
  verify(mockContext, times(0)).scheduleTasks(anyList());
  manager.onSourceTaskCompleted(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 2));
  verify(mockContext, times(1)).scheduleTasks(requestCaptor.capture());
  Assert.assertEquals(2, requestCaptor.getValue().size());
}
 

@Test(timeout = 5000)
public void testEdgeManagerPluginCtxGetVertexGroupName() throws TezException {
  EdgeManagerPluginDescriptor edgeManagerDescriptor =
    EdgeManagerPluginDescriptor.create(EdgeManagerForTest.class.getName());
  EdgeProperty edgeProp = EdgeProperty.create(edgeManagerDescriptor,
    DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("Out"),
    InputDescriptor.create("In"));
  Edge edge = new Edge(edgeProp, null, null);

  Vertex srcV = mock(Vertex.class), destV = mock(Vertex.class);
  String srcName = "srcV", destName = "destV";
  when(srcV.getName()).thenReturn(srcName);
  when(destV.getName()).thenReturn(destName);
  edge.setSourceVertex(srcV);
  edge.setDestinationVertex(destV);

  assertNull(edge.edgeManager.getContext().getVertexGroupName());

  String group = "group";
  when(destV.getGroupInputSpecList())
    .thenReturn(Arrays.asList(new GroupInputSpec(group, Arrays.asList("v1", "v3"), null)));
  assertNull(edge.edgeManager.getContext().getVertexGroupName());

  when(destV.getGroupInputSpecList())
    .thenReturn(Arrays.asList(new GroupInputSpec(group, Arrays.asList(srcName, "v3"), null)));
  assertEquals(group, edge.edgeManager.getContext().getVertexGroupName());
}
 

/**
 * This is a convenience method for the typical usage of this edge, and creates an instance of
 * {@link org.apache.tez.dag.api.EdgeProperty} which is likely to be used. </p>
 * If custom edge properties are required, the methods to get the relevant payloads should be
 * used. </p>
 * * In this case - DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED,
 * EdgeProperty.SchedulingType.SEQUENTIAL
 *
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultEdgeProperty() {
  EdgeProperty edgeProperty = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER,
      EdgeProperty.DataSourceType.PERSISTED, EdgeProperty.SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(
          getOutputClassName()).setUserPayload(getOutputPayload()),
      InputDescriptor.create(
          getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

VertexManagerPluginContext createVertexManagerContext(
    String mockSrcVertexId1, int numTasksSrcVertexId1,
    String mockSrcVertexId2, int numTasksSrcVertexId2,
    String mockSrcVertexId3, int numTasksSrcVertexId3,
    String mockManagedVertexId, int numTasksmockManagedVertexId,
    List<Integer> scheduledTasks,
    Map<String, EdgeManagerPlugin> newEdgeManagers) {
  HashMap<String, EdgeProperty> mockInputVertices =
      new HashMap<String, EdgeProperty>();
  EdgeProperty eProp1 = EdgeProperty.create(
      EdgeProperty.DataMovementType.SCATTER_GATHER,
      EdgeProperty.DataSourceType.PERSISTED,
      EdgeProperty.SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));
  EdgeProperty eProp2 = EdgeProperty.create(
      EdgeProperty.DataMovementType.SCATTER_GATHER,
      EdgeProperty.DataSourceType.PERSISTED,
      EdgeProperty.SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));
  EdgeProperty eProp3 = EdgeProperty.create(
      EdgeProperty.DataMovementType.BROADCAST,
      EdgeProperty.DataSourceType.PERSISTED,
      EdgeProperty.SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));

  mockInputVertices.put(mockSrcVertexId1, eProp1);
  mockInputVertices.put(mockSrcVertexId2, eProp2);
  mockInputVertices.put(mockSrcVertexId3, eProp3);

  final VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
  when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
  when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
  when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(numTasksSrcVertexId1);
  when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(numTasksSrcVertexId2);
  when(mockContext.getVertexNumTasks(mockSrcVertexId3)).thenReturn(numTasksSrcVertexId3);
  when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(numTasksmockManagedVertexId);
  doAnswer(new ScheduledTasksAnswer(scheduledTasks)).when(
      mockContext).scheduleTasks(anyList());
  doAnswer(new reconfigVertexAnswer(mockContext, mockManagedVertexId,
      newEdgeManagers)).when(mockContext).reconfigureVertex(
      anyInt(), any(VertexLocationHint.class), anyMap());
  return mockContext;
}
 

@Test (timeout=5000)
public void testBasicOneToOne() throws Exception {
  HashMap<String, EdgeProperty> mockInputVertices = 
      new HashMap<String, EdgeProperty>();
  String mockSrcVertexId1 = "Vertex1";
  EdgeProperty eProp1 = EdgeProperty.create(
      EdgeProperty.DataMovementType.ONE_TO_ONE,
      EdgeProperty.DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));
  
  String mockManagedVertexId = "Vertex";
  
  VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
  when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
  when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
  when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(3);
  when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(3);
  mockInputVertices.put(mockSrcVertexId1, eProp1);
  
  InputReadyVertexManager manager = new InputReadyVertexManager(mockContext);
  manager.initialize();
  verify(mockContext, times(1)).vertexReconfigurationPlanned();
  // source vertex configured
  manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
  verify(mockContext, times(1)).doneReconfiguringVertex();
  verify(mockContext, times(0)).scheduleTasks(requestCaptor.capture());
  manager.onVertexStarted(Collections.singletonList(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 0)));
  verify(mockContext, times(1)).scheduleTasks(requestCaptor.capture());
  Assert.assertEquals(1, requestCaptor.getValue().size());
  Assert.assertEquals(0, requestCaptor.getValue().get(0).getTaskIndex());
  Assert.assertEquals(mockSrcVertexId1, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedTask().getVertexName());
  Assert.assertEquals(0, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedTask().getTaskIndex());
  manager.onSourceTaskCompleted(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 1));
  verify(mockContext, times(2)).scheduleTasks(requestCaptor.capture());
  Assert.assertEquals(1, requestCaptor.getValue().size());
  Assert.assertEquals(1, requestCaptor.getValue().get(0).getTaskIndex());
  Assert.assertEquals(mockSrcVertexId1, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedTask().getVertexName());
  Assert.assertEquals(1, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedTask().getTaskIndex());
  manager.onSourceTaskCompleted(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 2));
  verify(mockContext, times(3)).scheduleTasks(requestCaptor.capture());
  Assert.assertEquals(1, requestCaptor.getValue().size());
  Assert.assertEquals(2, requestCaptor.getValue().get(0).getTaskIndex());
  Assert.assertEquals(mockSrcVertexId1, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedTask().getVertexName());
  Assert.assertEquals(2, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedTask().getTaskIndex());
}
 

@Test (timeout=5000)
public void testDelayedConfigureOneToOne() throws Exception {
  HashMap<String, EdgeProperty> mockInputVertices = 
      new HashMap<String, EdgeProperty>();
  String mockSrcVertexId1 = "Vertex1";
  EdgeProperty eProp1 = EdgeProperty.create(
      EdgeProperty.DataMovementType.ONE_TO_ONE,
      EdgeProperty.DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));
  
  String mockManagedVertexId = "Vertex";
  
  VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
  when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
  when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
  when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(3);
  when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(3);
  mockInputVertices.put(mockSrcVertexId1, eProp1);
  
  InputReadyVertexManager manager = new InputReadyVertexManager(mockContext);
  manager.initialize();
  verify(mockContext, times(1)).vertexReconfigurationPlanned();
  verify(mockContext, times(0)).scheduleTasks(requestCaptor.capture());
  // ok to have source task complete before anything else
  manager.onSourceTaskCompleted(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 1));
  // first own vertex started
  manager.onVertexStarted(Collections.singletonList(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 0)));
  // no scheduling as we are not configured yet
  verify(mockContext, times(0)).scheduleTasks(requestCaptor.capture());
  // then source vertex configured. now we start
  manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
  verify(mockContext, times(1)).doneReconfiguringVertex();
  
  verify(mockContext, times(2)).scheduleTasks(requestCaptor.capture());
  manager.onSourceTaskCompleted(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 2));
  verify(mockContext, times(3)).scheduleTasks(requestCaptor.capture());
  Assert.assertEquals(1, requestCaptor.getValue().size());
  Assert.assertEquals(2, requestCaptor.getValue().get(0).getTaskIndex());
  Assert.assertEquals(mockSrcVertexId1, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedTask().getVertexName());
  Assert.assertEquals(2, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedTask().getTaskIndex());
}
 

@Test(timeout = 5000)
public void testBasicVertexWithConcurrentInput() throws Exception {
  HashMap<String, EdgeProperty> mockInputVertices =
      new HashMap<String, EdgeProperty>();
  String mockSrcVertexId1 = "Vertex1";
  int srcVertex1Parallelism = 2;
  EdgeProperty eProp1 = EdgeProperty.create(
      EdgeManagerPluginDescriptor.create(SilentEdgeManager.class.getName()),
      EdgeProperty.DataSourceType.EPHEMERAL,
      EdgeProperty.SchedulingType.CONCURRENT,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));

  String mockSrcVertexId2 = "Vertex2";
  int srcVertex2Parallelism = 3;
  EdgeProperty eProp2 = EdgeProperty.create(
      EdgeManagerPluginDescriptor.create(SilentEdgeManager.class.getName()),
      EdgeProperty.DataSourceType.EPHEMERAL,
      EdgeProperty.SchedulingType.CONCURRENT,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));

  String mockManagedVertexId = "Vertex";
  int vertexParallelism = 2;

  VertexManagerWithConcurrentInput.ConcurrentInputVertexManagerConfigBuilder configurer =
      VertexManagerWithConcurrentInput.createConfigBuilder(null);
  VertexManagerPluginDescriptor pluginDesc = configurer.build();

  VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
  when(mockContext.getUserPayload()).thenReturn(pluginDesc.getUserPayload());
  when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
  when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
  when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(vertexParallelism);
  when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(srcVertex1Parallelism);
  when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(srcVertex2Parallelism);
  mockInputVertices.put(mockSrcVertexId1, eProp1);
  mockInputVertices.put(mockSrcVertexId2, eProp2);

  VertexManagerWithConcurrentInput manager = new VertexManagerWithConcurrentInput(mockContext);
  when(mockContext.getUserPayload()).thenReturn(pluginDesc.getUserPayload());
  manager.initialize();
  when(mockContext.getUserPayload()).thenReturn(pluginDesc.getUserPayload());

  // source vertex 1 configured
  manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
  verify(mockContext, times(0)).scheduleTasks(requestCaptor.capture());

  // source vertex 2 configured
  manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId2, VertexState.CONFIGURED));
  verify(mockContext, times(0)).scheduleTasks(requestCaptor.capture());

  // then own vertex started
  manager.onVertexStarted(Collections.singletonList(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 0)));
  verify(mockContext, times(1)).scheduleTasks(requestCaptor.capture());
  Assert.assertEquals(0, manager.completedUpstreamTasks);
}
 

@SuppressWarnings({ "rawtypes" })
@Test (timeout = 5000)
public void testCompositeEventHandling() throws TezException {
  EventHandler eventHandler = mock(EventHandler.class);
  EdgeProperty edgeProp = EdgeProperty.create(DataMovementType.SCATTER_GATHER,
      DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, mock(OutputDescriptor.class),
      mock(InputDescriptor.class));
  Edge edge = new Edge(edgeProp, eventHandler, new TezConfiguration());
  
  TezVertexID srcVertexID = createVertexID(1);
  TezVertexID destVertexID = createVertexID(2);
  LinkedHashMap<TezTaskID, Task> srcTasks = mockTasks(srcVertexID, 1);
  LinkedHashMap<TezTaskID, Task> destTasks = mockTasks(destVertexID, 5);
  
  TezTaskID srcTaskID = srcTasks.keySet().iterator().next();
  
  Vertex srcVertex = mockVertex("src", srcVertexID, srcTasks);
  Vertex destVertex = mockVertex("dest", destVertexID, destTasks);
  
  edge.setSourceVertex(srcVertex);
  edge.setDestinationVertex(destVertex);
  edge.initialize();
  
  TezTaskAttemptID srcTAID = createTAIDForTest(srcTaskID, 2); // Task0, Attempt 0
  
  EventMetaData srcMeta = new EventMetaData(EventProducerConsumerType.OUTPUT, "consumerVertex", "producerVertex", srcTAID);
  
  // Verification via a CompositeEvent
  CompositeDataMovementEvent cdmEvent = CompositeDataMovementEvent.create(0, destTasks.size(),
      ByteBuffer.wrap("bytes".getBytes()));
  cdmEvent.setVersion(2); // AttemptNum
  TezEvent tezEvent = new TezEvent(cdmEvent, srcMeta);
  // Event setup to look like it would after the Vertex is done with it.

  edge.sendTezEventToDestinationTasks(tezEvent);
  verifyEvents(srcTAID, destTasks);

  // Same Verification via regular DataMovementEvents
  // Reset the mock
  resetTaskMocks(destTasks.values());

  for (int i = 0 ; i < destTasks.size() ; i++) {
    DataMovementEvent dmEvent = DataMovementEvent.create(i, ByteBuffer.wrap("bytes".getBytes()));
    dmEvent.setVersion(2);
    tezEvent = new TezEvent(dmEvent, srcMeta);
    edge.sendTezEventToDestinationTasks(tezEvent);
  }
  verifyEvents(srcTAID, destTasks);
}
 

/**
 * Return EdgeProperty that connects two vertices.
 *
 * @param from
 * @param to
 * @return EdgeProperty
 * @throws IOException
 */
private EdgeProperty newEdge(TezOperator from, TezOperator to)
        throws IOException {
    TezEdgeDescriptor edge = to.inEdges.get(from.getOperatorKey());
    PhysicalPlan combinePlan = edge.combinePlan;

    InputDescriptor in = InputDescriptor.create(edge.inputClassName);
    OutputDescriptor out = OutputDescriptor.create(edge.outputClassName);

    Configuration conf = ConfigurationUtil.toConfiguration(pc.getProperties(), false);
    if (!combinePlan.isEmpty()) {
        addCombiner(combinePlan, to, conf);
    }

    List<POLocalRearrangeTez> lrs = PlanHelper.getPhysicalOperators(from.plan,
            POLocalRearrangeTez.class);

    for (POLocalRearrangeTez lr : lrs) {
        if (lr.getOutputKey().equals(to.getOperatorKey().toString())) {
            byte keyType = lr.getKeyType();
            setIntermediateOutputKeyValue(keyType, conf, to, lr.isConnectedToPackage());
            // In case of secondary key sort, main key type is the actual key type
            conf.set("pig.reduce.key.type", Byte.toString(lr.getMainKeyType()));
            break;
        }
    }

    conf.setIfUnset(TezRuntimeConfiguration.TEZ_RUNTIME_PARTITIONER_CLASS,
            MRPartitioner.class.getName());

    if (edge.getIntermediateOutputKeyClass() != null) {
        conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_KEY_CLASS,
                edge.getIntermediateOutputKeyClass());
    }

    if (edge.getIntermediateOutputValueClass() != null) {
        conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_VALUE_CLASS,
                edge.getIntermediateOutputValueClass());
    }

    if (edge.getIntermediateOutputKeyComparatorClass() != null) {
        conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_KEY_COMPARATOR_CLASS,
                edge.getIntermediateOutputKeyComparatorClass());
    }

    conf.setBoolean(MRConfiguration.MAPPER_NEW_API, true);
    conf.set("pig.pigContext", ObjectSerializer.serialize(pc));
    conf.set("udf.import.list",
            ObjectSerializer.serialize(PigContext.getPackageImportList()));

    if(to.isGlobalSort() || to.isLimitAfterSort()){
        conf.set("pig.sortOrder",
                ObjectSerializer.serialize(to.getSortOrder()));
    }

    if (edge.isUseSecondaryKey()) {
        conf.set("pig.secondarySortOrder",
                ObjectSerializer.serialize(edge.getSecondarySortOrder()));
        conf.set(org.apache.hadoop.mapreduce.MRJobConfig.PARTITIONER_CLASS_ATTR,
                SecondaryKeyPartitioner.class.getName());
        // These needs to be on the vertex as well for POShuffleTezLoad to pick it up.
        // Tez framework also expects this to be per vertex and not edge. IFile.java picks
        // up keyClass and valueClass from vertex config. TODO - check with Tez folks
        // In MR - job.setSortComparatorClass() or MRJobConfig.KEY_COMPARATOR
        conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_KEY_COMPARATOR_CLASS,
                PigSecondaryKeyComparator.class.getName());
        // In MR - job.setOutputKeyClass() or MRJobConfig.OUTPUT_KEY_CLASS
        conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_KEY_CLASS, NullableTuple.class.getName());
        setGroupingComparator(conf, PigSecondaryKeyGroupComparator.class.getName());
    }

    if (edge.partitionerClass != null) {
        conf.set(org.apache.hadoop.mapreduce.MRJobConfig.PARTITIONER_CLASS_ATTR,
                edge.partitionerClass.getName());
    }

    conf.set("udf.import.list",
            ObjectSerializer.serialize(PigContext.getPackageImportList()));

    MRToTezHelper.processMRSettings(conf, globalConf);

    String historyString = convertToHistoryText("", conf);
    in.setUserPayload(TezUtils.createUserPayloadFromConf(conf)).setHistoryText(historyString);
    out.setUserPayload(TezUtils.createUserPayloadFromConf(conf)).setHistoryText(historyString);

    if (edge.dataMovementType!=DataMovementType.BROADCAST && to.getEstimatedParallelism()!=-1 && (to.isGlobalSort()||to.isSkewedJoin())) {
        // Use custom edge
        return EdgeProperty.create((EdgeManagerPluginDescriptor)null,
                edge.dataSourceType, edge.schedulingType, out, in);
        }

    return EdgeProperty.create(edge.dataMovementType, edge.dataSourceType,
            edge.schedulingType, out, in);
}