Improve performance of LongObjectPagedHashMap#removeAndAdd and ObjectObjectPagedHashMap#removeAndAdd (#114280)

Author: iverase

server/src/main/java/org/elasticsearch/common/util/LongObjectPagedHashMap.java

@@ -63,6 +63,7 @@ public T get(long key) {
      * an insertion.
      */
     public T put(long key, T value) {
+        assert value != null : "Null values are not supported";
         if (size >= maxSize) {
             assert size == maxSize;
             grow();
@@ -94,9 +95,6 @@ public T remove(long key) {
     }
 
     private T set(long key, T value) {
-        if (value == null) {
-            throw new IllegalArgumentException("Null values are not supported");
-        }
         for (long i = slot(hash(key), mask);; i = nextSlot(i, mask)) {
             final T previous = values.getAndSet(i, value);
             if (previous == null) {
@@ -116,7 +114,7 @@ private T set(long key, T value) {
 
     @Override
     public Iterator<Cursor<T>> iterator() {
-        return new Iterator<Cursor<T>>() {
+        return new Iterator<>() {
 
             boolean cached;
             final Cursor<T> cursor;
@@ -181,9 +179,21 @@ protected boolean used(long bucket) {
     protected void removeAndAdd(long index) {
         final long key = keys.get(index);
         final T value = values.getAndSet(index, null);
-        --size;
-        final T removed = set(key, value);
-        assert removed == null;
+        reset(key, value);
+    }
+
+    private void reset(long key, T value) {
+        final ObjectArray<T> values = this.values;
+        final long mask = this.mask;
+        for (long i = slot(hash(key), mask);; i = nextSlot(i, mask)) {
+            final T previous = values.get(i);
+            if (previous == null) {
+                // slot was free
+                keys.set(i, key);
+                values.set(i, value);
+                break;
+            }
+        }
     }
 
     public static final class Cursor<T> {

server/src/main/java/org/elasticsearch/common/util/ObjectObjectPagedHashMap.java

@@ -67,6 +67,7 @@ public V get(K key) {
      * an insertion.
      */
     public V put(K key, V value) {
+        assert value != null : "Null values are not supported";
         if (size >= maxSize) {
             assert size == maxSize;
             grow();
@@ -100,7 +101,6 @@ public V remove(K key) {
 
     private V set(K key, int code, V value) {
         assert key.hashCode() == code;
-        assert value != null;
         assert size < maxSize;
         final long slot = slot(code, mask);
         for (long index = slot;; index = nextSlot(index, mask)) {
@@ -187,9 +187,22 @@ protected boolean used(long bucket) {
     protected void removeAndAdd(long index) {
         final K key = keys.get(index);
         final V value = values.getAndSet(index, null);
-        --size;
-        final V removed = set(key, key.hashCode(), value);
-        assert removed == null;
+        reset(key, value);
+    }
+
+    private void reset(K key, V value) {
+        final ObjectArray<V> values = this.values;
+        final long mask = this.mask;
+        final long slot = slot(key.hashCode(), mask);
+        for (long index = slot;; index = nextSlot(index, mask)) {
+            final V previous = values.get(index);
+            if (previous == null) {
+                // slot was free
+                values.set(index, value);
+                keys.set(index, key);
+                break;
+            }
+        }
     }
 
     public static final class Cursor<K, V> {