From b201535c3cafcaa010e1fbd6c72e69907234f491 Mon Sep 17 00:00:00 2001
From: Ivan Rakov <ivan.glukos@gmail.com>
Date: Wed, 30 Jun 2021 18:32:47 +0300
Subject: [PATCH] Fix edge cases reproduced on Guava CHM tests, closes
 https://github.com/OpenHFT/Chronicle-Map/issues/303

---
 .../net/openhft/chronicle/map/MapMethods.java |  9 ++++++++-
 .../chronicle/map/VanillaChronicleMap.java    | 20 +++++++++++++++----
 .../net/openhft/chronicle/map/GuavaTest.java  | 16 ---------------
 3 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/src/main/java/net/openhft/chronicle/map/MapMethods.java b/src/main/java/net/openhft/chronicle/map/MapMethods.java
index 511fdda9a..82363b2ec 100644
--- a/src/main/java/net/openhft/chronicle/map/MapMethods.java
+++ b/src/main/java/net/openhft/chronicle/map/MapMethods.java
@@ -18,6 +18,7 @@ package net.openhft.chronicle.map;
 
 import net.openhft.chronicle.hash.Data;
 
+import java.util.Objects;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
@@ -188,8 +189,13 @@ public interface MapMethods<K, V, R> {
                                  ReturnValue<V> returnValue) {
         if (tryReturnCurrentValueIfPresent(q, returnValue))
             return;
+
         // Key is absent
-        q.insert(q.absentEntry(), q.wrapValueAsData(mappingFunction.apply(q.queriedKey().get())));
+        final V applyResult = mappingFunction.apply(q.queriedKey().get());
+        if (applyResult == null)
+            return; // `null` is computed as a new value, no changes in the map are needed.
+
+        q.insert(q.absentEntry(), q.wrapValueAsData(applyResult));
         returnValue.returnValue(q.entry().value());
     }
 
@@ -415,6 +421,7 @@ public interface MapMethods<K, V, R> {
     default void merge(MapQueryContext<K, V, R> q, Data<V> value,
                        BiFunction<? super V, ? super V, ? extends V> remappingFunction,
                        ReturnValue<V> returnValue) {
+        Objects.requireNonNull(remappingFunction);
         q.updateLock().lock();
         Data<V> newValueData;
         MapEntry<K, V> entry = q.entry();
diff --git a/src/main/java/net/openhft/chronicle/map/VanillaChronicleMap.java b/src/main/java/net/openhft/chronicle/map/VanillaChronicleMap.java
index 334a6dc4a..d34c2ecbe 100644
--- a/src/main/java/net/openhft/chronicle/map/VanillaChronicleMap.java
+++ b/src/main/java/net/openhft/chronicle/map/VanillaChronicleMap.java
@@ -49,10 +49,7 @@ import java.io.IOException;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Type;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
@@ -823,6 +820,21 @@ public class VanillaChronicleMap<K, V, R>
         }
     }
 
+    @Override
+    public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        throwExceptionIfClosed();
+
+        Objects.requireNonNull(function);
+
+        forEachEntry(entry -> {
+            final V newValue = function.apply(entry.key().get(), entry.value().get());
+
+            // This is quite more simple than default ConcurrentMap#replaceAll implementation.
+            // Update lock for the entry is already acquired by #forEachEntry, so we are sure that the replace is atomic.
+            entry.context().replaceValue(entry, entry.context().wrapValueAsData(newValue));
+        });
+    }
+
     @Override
     public V replace(@NotNull final K key, @NotNull final V value) {
         throwExceptionIfClosed();
diff --git a/src/test/java/net/openhft/chronicle/map/GuavaTest.java b/src/test/java/net/openhft/chronicle/map/GuavaTest.java
index 97e5ad09e..5a442d4d7 100644
--- a/src/test/java/net/openhft/chronicle/map/GuavaTest.java
+++ b/src/test/java/net/openhft/chronicle/map/GuavaTest.java
@@ -44,22 +44,6 @@ public class GuavaTest extends TestCase {
         configureSuite(chmSuite);
         TestSuite chmTests = chmSuite.named("Guava tests of Chronicle Map").createTestSuite();
 
-        // TODO: remove after https://github.com/OpenHFT/Chronicle-Map/issues/303 is fixed.
-        {
-            Vector<TestSuite> collectionSizeSuites = Jvm.getValue(chmTests, "fTests");
-            assert collectionSizeSuites != null;
-
-            for (TestSuite sizeSuite : collectionSizeSuites) {
-                Vector<TestSuite> testers = Jvm.getValue(sizeSuite, "fTests");
-
-                assert testers != null;
-                testers.removeIf(tester ->
-                        tester.getName().contains("MapComputeIfAbsentTester") ||
-                                tester.getName().contains("MapMergeTester") ||
-                                tester.getName().contains("MapReplaceAllTester"));
-            }
-        }
-
         MapTestSuiteBuilder<String, String> backed = using(new BackedUpMapGenerator());
         configureSuite(backed);
         TestSuite backedTests = backed