diff --git a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfo.java b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfo.java
index 6593f4a1d..4c4c2a9d1 100644
--- a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfo.java
+++ b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfo.java
@@ -73,7 +73,7 @@ public class ShuffleTaskInfo {
     if (shuffleIdToTaskShuffleId.containsKey(shuffleId)) {
       String taskShuffleId = shuffleIdToTaskShuffleId.remove(shuffleId);
       taskShuffleIdToShuffleId.remove(taskShuffleId);
-      taskShuffleAttemptIdIndex.remove(shuffleId);
+      taskShuffleAttemptIdIndex.remove(taskShuffleId);
       taskShuffleAttemptIdToAttemptId.remove(taskShuffleId);
     }
   }
diff --git a/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfoSuitJ.java b/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfoSuitJ.java
index fd6325cf5..7acaa6405 100644
--- a/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfoSuitJ.java
+++ b/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfoSuitJ.java
@@ -47,10 +47,13 @@ public class ShuffleTaskInfoSuitJ {
     shuffleTaskInfo.removeExpiredShuffle(encodeShuffleId);
     int encodeShuffleIdNew = shuffleTaskInfo.getShuffleId("shuffleId");
     Assert.assertEquals(encodeShuffleIdNew, 2);
+
+    int encodeAttemptId211 = shuffleTaskInfo.getAttemptId("shuffleId", 1, "attempt1");
+    Assert.assertEquals(encodeAttemptId211, 0);
   }
 
   @Test
-  public void testRemoveNonExistShuffl() {
+  public void testRemoveNonExistShuffle() {
     ShuffleTaskInfo shuffleTaskInfo = new ShuffleTaskInfo();
     // remove none exist shuffle
     shuffleTaskInfo.removeExpiredShuffle(0);