ConcurrentModificationException in TaskAssignmentManager when partition count changes.

We need to make a copy of the keyset before deleting entries.

Exception in thread "main" java.util.ConcurrentModificationException
at java.util.HashMap$HashIterator.nextNode(HashMap.java:1429)
at java.util.HashMap$KeyIterator.next(HashMap.java:1453)
at java.util.Collections$UnmodifiableCollection$1.next(Collections.java:1102)
at org.apache.samza.container.grouper.task.TaskAssignmentManager.deleteTaskContainerMappings(TaskAssignmentManager.java:127)
at org.apache.samza.container.grouper.task.GroupByContainerCount.getPreviousContainers(GroupByContainerCount.java:158)
at org.apache.samza.container.grouper.task.GroupByContainerCount.balance(GroupByContainerCount.java:89)
at org.apache.samza.coordinator.JobCoordinator$.refreshJobModel(JobCoordinator.scala:254)
at org.apache.samza.coordinator.JobCoordinator$.jobModelGenerator$1(JobCoordinator.scala:186)
at org.apache.samza.coordinator.JobCoordinator$.initializeJobModel(JobCoordinator.scala:192)
at org.apache.samza.coordinator.JobCoordinator$.getJobCoordinator(JobCoordinator.scala:119)
at org.apache.samza.coordinator.JobCoordinator$.apply(JobCoordinator.scala:104)
at org.apache.samza.job.yarn.SamzaAppMaster$.main(SamzaAppMaster.scala:74)
at org.apache.samza.job.yarn.SamzaAppMaster.main(SamzaAppMaster.scala)

I have a unit test that reproduces the issue