From 523fe77013d27d1f4a5e659a3f23ae2a7ae047aa Mon Sep 17 00:00:00 2001
From: Simon Frankenberger <simon@fraho.eu>
Date: Thu, 17 Mar 2022 14:27:15 +0100
Subject: [PATCH] Iterating of the map should not throw an
 ConcurrentModificationException

Create a copy of the entries (guarded by a readLock) and iterate over that copy.
This fixes #10
---
 .../net/jodah/expiringmap/ExpiringMap.java    | 42 +++++++++-----
 .../jodah/expiringmap/ExpiringMapTest.java    |  5 +-
 .../net/jodah/expiringmap/issues/Issue10.java | 58 +++++++++++++++++++
 3 files changed, 90 insertions(+), 15 deletions(-)
 create mode 100644 src/test/java/net/jodah/expiringmap/issues/Issue10.java

diff --git a/src/main/java/net/jodah/expiringmap/ExpiringMap.java b/src/main/java/net/jodah/expiringmap/ExpiringMap.java
index 9edc5af..37912bb 100644
--- a/src/main/java/net/jodah/expiringmap/ExpiringMap.java
+++ b/src/main/java/net/jodah/expiringmap/ExpiringMap.java
@@ -4,6 +4,7 @@
 import java.util.AbstractCollection;
 import java.util.AbstractSet;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -50,8 +51,8 @@
  * Example usages:
  * 
  * <pre>
- * {@code 
- * Map<String, Integer> map = ExpiringMap.create(); 
+ * {@code
+ * Map<String, Integer> map = ExpiringMap.create();
  * Map<String, Integer> map = ExpiringMap.builder().expiration(30, TimeUnit.SECONDS).build();
  * Map<String, Connection> map = ExpiringMap.builder()
  *   .expiration(10, TimeUnit.MINUTES)
@@ -60,10 +61,10 @@
  *       return new Connection(address);
  *     }
  *   })
- *   .expirationListener(new ExpirationListener<String, Connection>() { 
- *     public void expired(String key, Connection connection) { 
+ *   .expirationListener(new ExpirationListener<String, Connection>() {
+ *     public void expired(String key, Connection connection) {
  *       connection.close();
- *     } 
+ *     }
  *   })
  *   .build();
  * }
@@ -367,9 +368,14 @@ public Iterator<ExpiringEntry<K, V>> valuesIterator() {
     }
 
     abstract class AbstractHashIterator {
-      private final Iterator<Map.Entry<K, ExpiringEntry<K, V>>> iterator = entrySet().iterator();
+      private final Iterator<Map.Entry<K, ExpiringEntry<K, V>>> iterator;
       private ExpiringEntry<K, V> next;
 
+      @SuppressWarnings({"unchecked", "rawtypes"})
+      AbstractHashIterator() {
+        iterator = (Iterator) Arrays.asList(entrySet().toArray(new Map.Entry[0])).iterator();
+      }
+
       public boolean hasNext() {
         return iterator.hasNext();
       }
@@ -385,19 +391,19 @@ public void remove() {
     }
 
     final class KeyIterator extends AbstractHashIterator implements Iterator<K> {
-      public final K next() {
+      public K next() {
         return getNext().key;
       }
     }
 
     final class ValueIterator extends AbstractHashIterator implements Iterator<V> {
-      public final V next() {
+      public V next() {
         return getNext().value;
       }
     }
 
     public final class EntryIterator extends AbstractHashIterator implements Iterator<Map.Entry<K, V>> {
-      public final Map.Entry<K, V> next() {
+      public Map.Entry<K, V> next() {
         return mapEntryFor(getNext());
       }
     }
@@ -879,8 +885,13 @@ public boolean contains(Object key) {
 
       @Override
       public Iterator<K> iterator() {
-        return (entries instanceof EntryLinkedHashMap) ? ((EntryLinkedHashMap<K, V>) entries).new KeyIterator()
-            : ((EntryTreeHashMap<K, V>) entries).new KeyIterator();
+        readLock.lock();
+        try {
+          return (entries instanceof EntryLinkedHashMap) ? ((EntryLinkedHashMap<K, V>) entries).new KeyIterator()
+                  : ((EntryTreeHashMap<K, V>) entries).new KeyIterator();
+        } finally {
+          readLock.unlock();
+        }
       }
 
       @Override
@@ -1195,8 +1206,13 @@ public boolean contains(Object value) {
 
       @Override
       public Iterator<V> iterator() {
-        return (entries instanceof EntryLinkedHashMap) ? ((EntryLinkedHashMap<K, V>) entries).new ValueIterator()
-            : ((EntryTreeHashMap<K, V>) entries).new ValueIterator();
+        readLock.lock();
+        try {
+          return (entries instanceof EntryLinkedHashMap) ? ((EntryLinkedHashMap<K, V>) entries).new ValueIterator()
+                  : ((EntryTreeHashMap<K, V>) entries).new ValueIterator();
+        } finally {
+          readLock.unlock();
+        }
       }
 
       @Override
diff --git a/src/test/java/net/jodah/expiringmap/ExpiringMapTest.java b/src/test/java/net/jodah/expiringmap/ExpiringMapTest.java
index eecb8a7..06b2c0b 100644
--- a/src/test/java/net/jodah/expiringmap/ExpiringMapTest.java
+++ b/src/test/java/net/jodah/expiringmap/ExpiringMapTest.java
@@ -441,10 +441,11 @@ public void testValues() {
   /**
    * Asserts that concurrent modification throws an exception.
    */
-  @Test(expectedExceptions = ConcurrentModificationException.class)
-  public void testValuesThrowsOnConcurrentModification() {
+  @Test
+  public void testValuesNotThrowsOnConcurrentModification() {
     ExpiringMap<String, String> map = ExpiringMap.create();
     map.put("a", "a");
+    map.put("b", "b");
 
     Iterator<String> valuesIterator = map.values().iterator();
     valuesIterator.next();
diff --git a/src/test/java/net/jodah/expiringmap/issues/Issue10.java b/src/test/java/net/jodah/expiringmap/issues/Issue10.java
new file mode 100644
index 0000000..e65e6e0
--- /dev/null
+++ b/src/test/java/net/jodah/expiringmap/issues/Issue10.java
@@ -0,0 +1,58 @@
+package net.jodah.expiringmap.issues;
+
+import net.jodah.expiringmap.ExpirationPolicy;
+import net.jodah.expiringmap.ExpiringMap;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.util.Iterator;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Do not throw ConcurrentModificationException when using an Iterator
+ */
+@Test
+public class Issue10 {
+    public void testEntrySet() {
+        Map<Integer, Integer> map = ExpiringMap.builder().expiration(30, TimeUnit.SECONDS).expirationPolicy(ExpirationPolicy.ACCESSED).build();
+        map.put(1, 1);
+        map.put(2, 2);
+        Iterator<Map.Entry<Integer, Integer>> iterator = map.entrySet().iterator();
+        Assert.assertEquals((Integer) 1, iterator.next().getKey());
+        map.put(3, 3);
+        Assert.assertEquals((Integer) 2, iterator.next().getKey());
+        // we don't require the iterator to be updated with the new '3' entry and we neither expect a ConcurrentModificationException
+        if (iterator.hasNext()) {
+            Assert.assertEquals((Integer) 3, iterator.next().getKey());
+        }
+    }
+
+    public void testValues() {
+        Map<Integer, Integer> map = ExpiringMap.builder().expiration(30, TimeUnit.SECONDS).expirationPolicy(ExpirationPolicy.ACCESSED).build();
+        map.put(1, 1);
+        map.put(2, 2);
+        Iterator<Integer> iterator = map.values().iterator();
+        Assert.assertEquals((Integer) 1, iterator.next());
+        map.put(3, 3);
+        Assert.assertEquals((Integer) 2, iterator.next());
+        // we don't require the iterator to be updated with the new '3' entry and we neither expect a ConcurrentModificationException
+        if (iterator.hasNext()) {
+            Assert.assertEquals((Integer) 3, iterator.next());
+        }
+    }
+
+    public void testKeySet() {
+        Map<Integer, Integer> map = ExpiringMap.builder().expiration(30, TimeUnit.SECONDS).expirationPolicy(ExpirationPolicy.ACCESSED).build();
+        map.put(1, 1);
+        map.put(2, 2);
+        Iterator<Integer> iterator = map.keySet().iterator();
+        Assert.assertEquals((Integer) 1, iterator.next());
+        map.put(3, 3);
+        Assert.assertEquals((Integer) 2, iterator.next());
+        // we don't require the iterator to be updated with the new '3' entry and we neither expect a ConcurrentModificationException
+        if (iterator.hasNext()) {
+            Assert.assertEquals((Integer) 3, iterator.next());
+        }
+    }
+}