Fix ConcurrentModificationException when we traverse registeredTasks #…

…664 (#665) * [bugfix] ConcurrentModificationException when we traverse registeredTasks #664 * remove unused import
tony-framework · Apr 28, 2022 · 6e0985b · 6e0985b
1 parent 63fe90f
commit 6e0985b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/build.gradle b/build.gradle
@@ -92,7 +92,7 @@ ext.deps = [
 
 allprojects {
   group = "com.linkedin.tony"
-  project.version = "0.5.1"
+  project.version = "0.5.2"
 }
 
 task sourcesJar(type: Jar) {

diff --git a/tony-core/src/main/java/com/linkedin/tony/TonySession.java b/tony-core/src/main/java/com/linkedin/tony/TonySession.java
@@ -14,7 +14,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -67,7 +66,7 @@ public class TonySession {
   // go straight to the cleaning phase.
   private boolean trainingFinished = false;
 
-  private Set<String> registeredTasks = new HashSet<>();
+  private Set<String> registeredTasks = ConcurrentHashMap.newKeySet();
 
   private int numExpectedTasks = 0;
 
@@ -677,7 +676,7 @@ public void addRegisteredTask(String taskId) {
   }
 
   public void resetRegisteredTasks() {
-    registeredTasks = new HashSet<>();
+    registeredTasks = ConcurrentHashMap.newKeySet();
   }
 
   public int getNumRegisteredTasks() {