From 228290b2fcf327b3ed1dc36092aacf49ac735491 Mon Sep 17 00:00:00 2001
From: Simon Frankenberger <simon@fraho.eu>
Date: Tue, 5 Dec 2023 11:14:39 +0100
Subject: [PATCH] Use the readLock when creating an Iterator for the Map.

This is an addition to the fix implemented with #78 which did not cover all cases.
---
 .../net/jodah/expiringmap/ExpiringMap.java    |  9 +++--
 .../net/jodah/expiringmap/issues/Issue10.java | 36 +++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/main/java/net/jodah/expiringmap/ExpiringMap.java b/src/main/java/net/jodah/expiringmap/ExpiringMap.java
index 1a4131c..26dc4e9 100644
--- a/src/main/java/net/jodah/expiringmap/ExpiringMap.java
+++ b/src/main/java/net/jodah/expiringmap/ExpiringMap.java
@@ -714,8 +714,13 @@ public boolean contains(Object entry) {
 
       @Override
       public Iterator<Map.Entry<K, V>> iterator() {
-        return (entries instanceof EntryLinkedHashMap) ? ((EntryLinkedHashMap<K, V>) entries).new EntryIterator()
-            : ((EntryTreeHashMap<K, V>) entries).new EntryIterator();
+        readLock.lock();
+        try {
+          return (entries instanceof EntryLinkedHashMap) ? ((EntryLinkedHashMap<K, V>) entries).new EntryIterator()
+                  : ((EntryTreeHashMap<K, V>) entries).new EntryIterator();
+        } finally {
+          readLock.unlock();
+        }
       }
 
       @Override
diff --git a/src/test/java/net/jodah/expiringmap/issues/Issue10.java b/src/test/java/net/jodah/expiringmap/issues/Issue10.java
index e65e6e0..9f64bf0 100644
--- a/src/test/java/net/jodah/expiringmap/issues/Issue10.java
+++ b/src/test/java/net/jodah/expiringmap/issues/Issue10.java
@@ -5,9 +5,15 @@
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import java.util.ArrayList;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * Do not throw ConcurrentModificationException when using an Iterator
@@ -55,4 +61,34 @@ public void testKeySet() {
             Assert.assertEquals((Integer) 3, iterator.next());
         }
     }
+
+    public void testParallelPutRemove() {
+        ExpiringMap<Integer, String> testee = ExpiringMap.create();
+        ExecutorService service = Executors.newFixedThreadPool(2);
+
+        Runnable adder = () -> {
+            ThreadLocalRandom rnd = ThreadLocalRandom.current();
+            AtomicInteger counter = new AtomicInteger(0);
+            for (int j = 0; j < 5; j++) {
+                for (int i = 0; i < 100_000; i++) {
+                    if (rnd.nextBoolean()) {
+                        testee.put(counter.incrementAndGet(), "bar");
+                    }
+                }
+            }
+        };
+
+        Runnable remover = () -> {
+            while (!service.isTerminated()) {
+                List<Integer> entriesToDelete = new ArrayList<>();
+                for (Map.Entry<Integer, String> e : testee.entrySet()) {
+                    entriesToDelete.add(e.getKey());
+                }
+                entriesToDelete.forEach(testee.keySet()::remove);
+            }
+        };
+        service.submit(adder);
+        service.shutdown(); // schedule shutdown, let the running tasks finish
+        remover.run(); // run synchronous, waits for the adder threads to finish
+    }
 }