@SuppressWarnings("unchecked")
 public void sendTezEventToSourceTasks(TezEvent tezEvent) {
   if (!bufferEvents.get()) {
     switch (tezEvent.getEventType()) {
       case INPUT_READ_ERROR_EVENT:
         InputReadErrorEvent event = (InputReadErrorEvent) tezEvent.getEvent();
         TezTaskAttemptID destAttemptId = tezEvent.getSourceInfo().getTaskAttemptID();
         int destTaskIndex = destAttemptId.getTaskID().getId();
         int srcTaskIndex = edgeManager.routeEventToSourceTasks(destTaskIndex, event);
         int numConsumers =
             edgeManager.getDestinationConsumerTaskNumber(
                 srcTaskIndex, destinationVertex.getTotalTasks());
         TezTaskID srcTaskId = sourceVertex.getTask(srcTaskIndex).getTaskId();
         int taskAttemptIndex = event.getVersion();
         TezTaskAttemptID srcTaskAttemptId = new TezTaskAttemptID(srcTaskId, taskAttemptIndex);
         eventHandler.handle(
             new TaskAttemptEventOutputFailed(srcTaskAttemptId, tezEvent, numConsumers));
         break;
       default:
         throw new TezUncheckedException("Unhandled tez event type: " + tezEvent.getEventType());
     }
   } else {
     sourceEventBuffer.add(tezEvent);
   }
 }
 public void sendTezEventToDestinationTasks(TezEvent tezEvent) {
   if (!bufferEvents.get()) {
     List<Integer> destTaskIndices = new ArrayList<Integer>();
     boolean isDataMovementEvent = true;
     switch (tezEvent.getEventType()) {
       case INPUT_FAILED_EVENT:
         isDataMovementEvent = false;
       case DATA_MOVEMENT_EVENT:
         Event event = tezEvent.getEvent();
         TezTaskAttemptID sourceAttemptId = tezEvent.getSourceInfo().getTaskAttemptID();
         int sourceTaskIndex = sourceAttemptId.getTaskID().getId();
         if (isDataMovementEvent) {
           edgeManager.routeEventToDestinationTasks(
               (DataMovementEvent) event,
               sourceTaskIndex,
               destinationVertex.getTotalTasks(),
               destTaskIndices);
         } else {
           edgeManager.routeEventToDestinationTasks(
               (InputFailedEvent) event,
               sourceTaskIndex,
               destinationVertex.getTotalTasks(),
               destTaskIndices);
         }
         for (Integer destTaskIndex : destTaskIndices) {
           EventMetaData destMeta =
               new EventMetaData(
                   EventProducerConsumerType.INPUT,
                   destinationVertex.getName(),
                   sourceVertex.getName(),
                   null); // will be filled by Task when sending the event. Is it needed?
           if (isDataMovementEvent) {
             destMeta.setIndex(((DataMovementEvent) event).getTargetIndex());
           } else {
             destMeta.setIndex(((InputFailedEvent) event).getTargetIndex());
           }
           tezEvent.setDestinationInfo(destMeta);
           TezTaskID destTaskId = destinationVertex.getTask(destTaskIndex).getTaskId();
           sendEventToTask(destTaskId, tezEvent);
         }
         break;
       default:
         throw new TezUncheckedException("Unhandled tez event type: " + tezEvent.getEventType());
     }
   } else {
     destinationEventBuffer.add(tezEvent);
   }
 }
    @Override
    public TezHeartbeatResponse heartbeat(TezHeartbeatRequest request)
        throws IOException, TezException {
      // Keep-alive information. The client should be informed and will have to take care of
      // re-submitting the work.
      // Some parts of fault tolerance go here.

      // This also provides completion information, and a possible notification when task actually
      // starts running (first heartbeat)

      if (LOG.isDebugEnabled()) {
        LOG.debug("Received heartbeat from container, request=" + request);
      }

      // Incoming events can be ignored until the point when shuffle needs to be handled, instead of
      // just scans.
      TezHeartbeatResponse response = new TezHeartbeatResponse();

      response.setLastRequestId(request.getRequestId());
      // Assuming TaskAttemptId and FragmentIdentifierString are the same. Verify this.
      TezTaskAttemptID taskAttemptId = request.getCurrentTaskAttemptID();
      String taskAttemptIdString = taskAttemptId.toString();

      updateHeartbeatInfo(taskAttemptIdString);

      List<TezEvent> tezEvents = null;
      PendingEventData pendingEventData = pendingEvents.remove(taskAttemptIdString);
      if (pendingEventData == null) {
        tezEvents = Collections.emptyList();

        // If this heartbeat was not from a pending event and it's not in our list of registered
        // tasks,
        if (!registeredTasks.containsKey(taskAttemptIdString)) {
          LOG.info("Unexpected heartbeat from " + taskAttemptIdString);
          response.setShouldDie(); // Do any of the other fields need to be set?
          return response;
        }
      } else {
        tezEvents = pendingEventData.tezEvents;
        // Tasks removed from the pending list should then be added to the registered list.
        registeredTasks.put(taskAttemptIdString, pendingEventData.heartbeatInfo);
      }

      response.setLastRequestId(request.getRequestId());
      // Irrelevant from eventIds. This can be tracked in the AM itself, instead of polluting the
      // task.
      // Also since we have all the MRInput events here - they'll all be sent in together.
      response.setNextFromEventId(0); // Irrelevant. See comment above.
      response.setNextPreRoutedEventId(0); // Irrelevant. See comment above.
      response.setEvents(tezEvents);

      List<TezEvent> inEvents = request.getEvents();
      if (LOG.isDebugEnabled()) {
        LOG.debug(
            "Heartbeat from "
                + taskAttemptIdString
                + " events: "
                + (inEvents != null ? inEvents.size() : -1));
      }
      for (TezEvent tezEvent : ListUtils.emptyIfNull(inEvents)) {
        EventType eventType = tezEvent.getEventType();
        switch (eventType) {
          case TASK_ATTEMPT_COMPLETED_EVENT:
            LOG.debug("Task completed event for " + taskAttemptIdString);
            registeredTasks.remove(taskAttemptIdString);
            break;
          case TASK_ATTEMPT_FAILED_EVENT:
            LOG.debug("Task failed event for " + taskAttemptIdString);
            registeredTasks.remove(taskAttemptIdString);
            break;
          case TASK_STATUS_UPDATE_EVENT:
            // If we want to handle counters
            LOG.debug("Task update event for " + taskAttemptIdString);
            break;
          default:
            LOG.warn("Unhandled event type " + eventType);
            break;
        }
      }

      // Pass the request on to the responder
      try {
        if (responder != null) {
          responder.heartbeat(request);
        }
      } catch (Exception err) {
        LOG.error("Error during responder execution", err);
      }

      return response;
    }