TEZ-4061: InputAttemptIdentifier and CompositeInputAttemptIdentifier cannot be compared for equality

tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/CompositeInputAttemptIdentifier.java

@@ -50,6 +50,12 @@ public InputAttemptIdentifier expand(int inputIdentifierOffset) {
     return new InputAttemptIdentifier(getInputIdentifier() + inputIdentifierOffset, getAttemptNumber(), getPathComponent(), isShared(), getFetchTypeInfo(), getSpillEventId());
   }
 
+  public boolean include(int thatInputIdentifier, int thatAttemptNumber) {
+    return
+        super.getInputIdentifier() <= thatInputIdentifier && thatInputIdentifier < (super.getInputIdentifier() + inputIdentifierCount) &&
+        super.getAttemptNumber() == thatAttemptNumber;
+  }
+
   // PathComponent & shared does not need to be part of the hashCode and equals computation.
   @Override
   public int hashCode() {
@@ -63,6 +69,6 @@ public boolean equals(Object obj) {
 
   @Override
   public String toString() {
-    return super.toString();
+    return super.toString() + ", count=" + inputIdentifierCount;
   }
 }

tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/InputAttemptIdentifier.java

@@ -108,6 +108,10 @@ public boolean canRetrieveInputInChunks() {
         (fetchTypeInfo == SPILL_INFO.FINAL_UPDATE.ordinal());
   }
 
+  public boolean include(int thatInputIdentifier, int thatAttemptNumber) {
+    return this.inputIdentifier == thatInputIdentifier && this.attemptNumber == thatAttemptNumber;
+  }
+
   // PathComponent & shared does not need to be part of the hashCode and equals computation.
   @Override
   public int hashCode() {
@@ -139,6 +143,6 @@ public boolean equals(Object obj) {
   public String toString() {
     return "InputAttemptIdentifier [inputIdentifier=" + inputIdentifier
         + ", attemptNumber=" + attemptNumber + ", pathComponent="
-        + pathComponent + ", spillType=" + fetchTypeInfo + ", spillId=" + spillEventId  +"]";
+        + pathComponent + ", spillType=" + fetchTypeInfo + ", spillId=" + spillEventId + "]";
   }
 }

tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/ShuffleInputEventHandlerImpl.java

@@ -282,6 +282,7 @@ private void processShufflePayload(DataMovementEventPayloadProto shufflePayload,
 
   private void processInputFailedEvent(InputFailedEvent ife) {
     InputAttemptIdentifier srcAttemptIdentifier = new InputAttemptIdentifier(ife.getTargetIndex(), ife.getVersion());
+    LOG.info("Marking obsolete input: " + inputContext.getSourceVertexName() + " " + srcAttemptIdentifier);
     shuffleManager.obsoleteKnownInput(srcAttemptIdentifier);
   }
 

tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/ShuffleManager.java

@@ -583,8 +583,15 @@ Fetcher constructFetcherForHost(InputHost inputHost, Configuration conf) {
       } else {
           alreadyCompleted = completedInputSet.get(input.getInputIdentifier());
       }
-      // Avoid adding attempts which have already completed or have been marked as OBSOLETE
-      if (alreadyCompleted || obsoletedInputs.contains(input)) {
+
+      // Avoid adding attempts which have already completed
+      if (alreadyCompleted) {
+        inputIter.remove();
+        continue;
+      }
+      // Avoid adding attempts which have been marked as OBSOLETE
+      if (isObsoleteInputAttemptIdentifier(input)) {
+        LOG.info("Skipping obsolete input: " + input);
         inputIter.remove();
         continue;
       }
@@ -949,10 +956,14 @@ public void fetchFailed(String host,
     // TODO NEWTEZ. Implement logic to report fetch failures after a threshold.
     // For now, reporting immediately.
     InputAttemptIdentifier srcAttemptIdentifier = inputAttemptFetchFailure.getInputAttemptIdentifier();
+    if (isObsoleteInputAttemptIdentifier(srcAttemptIdentifier)) {
+      LOG.info("Do not report obsolete input: " + srcAttemptIdentifier);
+      return;
+    }
     LOG.info(
-        "{}: Fetch failed for src: {} InputIdentifier: {}, connectFailed: {}, "
+        "{}: Fetch failed for InputIdentifier: {}, connectFailed: {}, "
             + "local fetch: {}, remote fetch failure reported as local failure: {})",
-        sourceDestNameTrimmed, srcAttemptIdentifier, srcAttemptIdentifier, connectFailed,
+        sourceDestNameTrimmed, srcAttemptIdentifier, connectFailed,
         inputAttemptFetchFailure.isLocalFetch(), inputAttemptFetchFailure.isDiskErrorAtSource());
     failedShufflesCounter.increment(1);
     inputContext.notifyProgress();
@@ -984,6 +995,22 @@ public void fetchFailed(String host,
       }
     }
   }
+
+  private boolean isObsoleteInputAttemptIdentifier(InputAttemptIdentifier input) {
+    if (input == null) {
+      return false;
+    }
+    InputAttemptIdentifier obsoleteInput;
+    Iterator<InputAttemptIdentifier> obsoleteInputsIter = obsoletedInputs.iterator();
+    while (obsoleteInputsIter.hasNext()) {
+      obsoleteInput = obsoleteInputsIter.next();
+      if (input.include(obsoleteInput.getInputIdentifier(), obsoleteInput.getAttemptNumber())) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   /////////////////// End of Methods from FetcherCallbackHandler
 
   public void shutdown() throws InterruptedException {

tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/ShuffleScheduler.java

@@ -1175,7 +1175,19 @@ private synchronized boolean inputShouldBeConsumed(InputAttemptIdentifier id) {
     } else {
       isInputFinished = isInputFinished(id.getInputIdentifier());
     }
-    return !obsoleteInputs.contains(id) && !isInputFinished;
+    return !isObsoleteInputAttemptIdentifier(id) && !isInputFinished;
+  }
+
+  private boolean isObsoleteInputAttemptIdentifier(InputAttemptIdentifier input) {
+    InputAttemptIdentifier obsoleteInput;
+    Iterator<InputAttemptIdentifier> obsoleteInputsIter = obsoleteInputs.iterator();
+    while (obsoleteInputsIter.hasNext()) {
+      obsoleteInput = obsoleteInputsIter.next();
+      if (input.include(obsoleteInput.getInputIdentifier(), obsoleteInput.getAttemptNumber())) {
+        return true;
+      }
+    }
+    return false;
   }
 
   public synchronized List<InputAttemptIdentifier> getMapsForHost(MapHost host) {