Fix ConcurrentModificationException when we traverse registeredTasks #664 (#665)

* [bugfix] ConcurrentModificationException when we traverse registeredTasks #664

* remove unused import

Author: Keqiu Hu

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) {

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() {