More performance improvements

benchmarks/src/jmh/java/io/micronaut/core/CopyOnWriteMapBenchmark.java

@@ -0,0 +1,83 @@
+package io.micronaut.core;
+
+import io.micronaut.core.util.CopyOnWriteMap;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+
+public class CopyOnWriteMapBenchmark {
+    public static void main(String[] args) throws RunnerException {
+        Options opt = new OptionsBuilder()
+            .include(CopyOnWriteMapBenchmark.class.getName() + ".*")
+            .warmupIterations(3)
+            .measurementIterations(5)
+            .mode(Mode.AverageTime)
+            .timeUnit(TimeUnit.NANOSECONDS)
+            .forks(1)
+            //.addProfiler(LinuxPerfAsmProfiler.class)
+            .build();
+
+        new Runner(opt).run();
+    }
+
+    @Benchmark
+    public String get(S s) {
+        return s.map.get("foo");
+    }
+
+    @Benchmark
+    public String computeIfAbsent(S s) {
+        return s.map.computeIfAbsent("fizz", s.ciaUpdate);
+    }
+
+    @Benchmark
+    public String getWithCheck(S s) {
+        String v = s.map.get("fizz");
+        if (v == null) {
+            return s.map.computeIfAbsent("fizz", s.ciaUpdate);
+        } else {
+            return v;
+        }
+    }
+
+    @State(Scope.Thread)
+    public static class S {
+        @Param({"CHM", "COW"})
+        Type type;
+        @Param({"1", "2", "5", "10"})
+        int load;
+        private Map<String, String> map;
+        private Function<String, String> ciaUpdate;
+
+        @Setup
+        public void setUp() {
+            map = switch (type) {
+                case CHM -> new ConcurrentHashMap<>(16, 0.75f, 1);
+                case COW -> new CopyOnWriteMap<>(16);
+            };
+            // doesn't really stress the collision avoidance algorithm but oh well
+            map.put("foo", "bar");
+            for (int i = 0; i < load; i++) {
+                map.put("f" + i, "b" + i);
+            }
+            ciaUpdate = m -> "buzz" + m;
+        }
+    }
+
+    public enum Type {
+        CHM,
+        COW,
+    }
+}

core/src/main/java/io/micronaut/core/util/CopyOnWriteMap.java

@@ -0,0 +1,281 @@
+package io.micronaut.core.util;
+
+import io.micronaut.core.annotation.Internal;
+import io.micronaut.core.annotation.NonNull;
+
+import java.util.AbstractMap;
+import java.util.AbstractSet;
+import java.util.BitSet;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+
+/**
+ * Thread-safe map that is optimized for reads. Uses a normal {@link HashMap} that is copied on
+ * update operations.
+ *
+ * @param <K> The key type
+ * @param <V> The value type
+ */
+@Internal
+public final class CopyOnWriteMap<K, V> extends AbstractMap<K, V> implements ConcurrentMap<K, V> {
+    /**
+     * Empty {@link HashMap} to avoid polymorphism.
+     */
+    @SuppressWarnings("rawtypes")
+    private static final Map EMPTY = new HashMap();
+    /**
+     * How many items to evict at a time, to make eviction a bit more efficient.
+     */
+    static final int EVICTION_BATCH = 16;
+
+    private final int maxSizeWithEvictionMargin;
+    @SuppressWarnings("unchecked")
+    private volatile Map<? extends K, ? extends V> actual = EMPTY;
+
+    public CopyOnWriteMap(int maxSize) {
+        int maxSizeWithEvictionMargin = maxSize + EVICTION_BATCH;
+        if (maxSizeWithEvictionMargin < 0) {
+            maxSizeWithEvictionMargin = Integer.MAX_VALUE;
+        }
+        this.maxSizeWithEvictionMargin = maxSizeWithEvictionMargin;
+    }
+
+    @NonNull
+    @Override
+    public Set<Entry<K, V>> entrySet() {
+        return new EntrySet();
+    }
+
+    @Override
+    public V get(Object key) {
+        return actual.get(key);
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public V getOrDefault(Object key, V defaultValue) {
+        return ((Map<K, V>) actual).getOrDefault(key, defaultValue);
+    }
+
+    @Override
+    public boolean containsKey(Object key) {
+        return actual.containsKey(key);
+    }
+
+    @Override
+    public boolean containsValue(Object value) {
+        return actual.containsValue(value);
+    }
+
+    @Override
+    public int size() {
+        return actual.size();
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public synchronized void clear() {
+        actual = EMPTY;
+    }
+
+    @Override
+    public void putAll(Map<? extends K, ? extends V> m) {
+        update(map -> {
+            map.putAll(m);
+            return null;
+        });
+    }
+
+    @Override
+    public V remove(Object key) {
+        return update(m -> m.remove(key));
+    }
+
+    @Override
+    public int hashCode() {
+        return actual.hashCode();
+    }
+
+    @SuppressWarnings("EqualsWhichDoesntCheckParameterClass")
+    @Override
+    public boolean equals(Object o) {
+        return actual.equals(o);
+    }
+
+    @Override
+    public String toString() {
+        return actual.toString();
+    }
+
+    @Override
+    public void forEach(BiConsumer<? super K, ? super V> action) {
+        actual.forEach(action);
+    }
+
+    private synchronized <R> R update(Function<Map<K, V>, R> updater) {
+        Map<K, V> next = new HashMap<>(actual);
+        R ret = updater.apply(next);
+        int newSize = next.size();
+        if (newSize >= maxSizeWithEvictionMargin) {
+            // select some indices in the map to remove at random
+            BitSet toRemove = new BitSet(newSize);
+            for (int i = 0; i < EVICTION_BATCH; i++) {
+                setUnset(toRemove, ThreadLocalRandom.current().nextInt(newSize - i));
+            }
+            // iterate over the map and remove those indices
+            Iterator<?> iterator = next.entrySet().iterator();
+            for (int i = 0; i < newSize; i++) {
+                iterator.next();
+                if (toRemove.get(i)) {
+                    iterator.remove();
+                }
+            }
+        }
+        actual = next;
+        return ret;
+    }
+
+    /**
+     * Set the bit at {@code index}, with the index only counting unset bits. e.g. setting index 0
+     * when the first bit of the {@link BitSet} is already set would set the second bit (the first
+     * unset bit).
+     *
+     * @param set   The bit set to modify
+     * @param index The index of the bit to set
+     */
+    static void setUnset(BitSet set, int index) {
+        int i = 0;
+        while (true) {
+            int nextI = set.nextSetBit(i);
+            if (nextI == -1 || nextI > index) {
+                break;
+            }
+            i = nextI + 1;
+            index++;
+        }
+        set.set(index);
+    }
+
+    @Override
+    public V put(K key, V value) {
+        return update(m -> m.put(key, value));
+    }
+
+    @Override
+    public boolean remove(@NonNull Object key, Object value) {
+        return update(m -> m.remove(key, value));
+    }
+
+    @Override
+    public boolean replace(@NonNull K key, @NonNull V oldValue, @NonNull V newValue) {
+        return update(m -> m.replace(key, oldValue, newValue));
+    }
+
+    @Override
+    public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        update(m -> {
+            m.replaceAll(function);
+            return null;
+        });
+    }
+
+    @Override
+    public V computeIfAbsent(K key, @NonNull Function<? super K, ? extends V> mappingFunction) {
+        V present = get(key);
+        if (present != null) {
+            // fast path without sync
+            return present;
+        } else {
+            return update(m -> m.computeIfAbsent(key, mappingFunction));
+        }
+    }
+
+    @Override
+    public V computeIfPresent(K key, @NonNull BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
+        return update(m -> m.computeIfPresent(key, remappingFunction));
+    }
+
+    @Override
+    public V compute(K key, @NonNull BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
+        return update(m -> m.compute(key, remappingFunction));
+    }
+
+    @Override
+    public V merge(K key, @NonNull V value, @NonNull BiFunction<? super V, ? super V, ? extends V> remappingFunction) {
+        return update(m -> m.merge(key, value, remappingFunction));
+    }
+
+    @Override
+    public V putIfAbsent(@NonNull K key, V value) {
+        return update(m -> m.putIfAbsent(key, value));
+    }
+
+    @Override
+    public V replace(@NonNull K key, @NonNull V value) {
+        return update(m -> m.replace(key, value));
+    }
+
+    private class EntrySetIterator implements Iterator<Entry<K, V>> {
+        final Iterator<? extends Entry<? extends K, ? extends V>> itr = actual.entrySet().iterator();
+        K lastKey;
+
+        @Override
+        public boolean hasNext() {
+            return itr.hasNext();
+        }
+
+        @Override
+        public Entry<K, V> next() {
+            Entry<? extends K, ? extends V> e = itr.next();
+            lastKey = e.getKey();
+            return new EntryImpl(e);
+        }
+
+        @Override
+        public void remove() {
+            CopyOnWriteMap.this.remove(lastKey);
+        }
+    }
+
+    private class EntryImpl implements Entry<K, V> {
+        private final Entry<? extends K, ? extends V> entry;
+
+        public EntryImpl(Entry<? extends K, ? extends V> entry) {
+            this.entry = entry;
+        }
+
+        @Override
+        public K getKey() {
+            return entry.getKey();
+        }
+
+        @Override
+        public V getValue() {
+            return entry.getValue();
+        }
+
+        @Override
+        public V setValue(V value) {
+            return put(entry.getKey(), value);
+        }
+    }
+
+    private class EntrySet extends AbstractSet<Entry<K, V>> {
+        @Override
+        public Iterator<Entry<K, V>> iterator() {
+            return new EntrySetIterator();
+        }
+
+        @Override
+        public int size() {
+            return actual.size();
+        }
+    }
+}