From c009ef05052e85f562ebed23081647ade23fbe04 Mon Sep 17 00:00:00 2001
From: themode <themode@outlook.fr>
Date: Wed, 4 May 2022 18:24:15 +0200
Subject: [PATCH] Avoid slow synchronized path during reference building

---
 .../server/snapshot/SnapshotUpdaterImpl.java  | 25 +++++++++++++------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/main/java/net/minestom/server/snapshot/SnapshotUpdaterImpl.java b/src/main/java/net/minestom/server/snapshot/SnapshotUpdaterImpl.java
index ed3758841..63b4a412c 100644
--- a/src/main/java/net/minestom/server/snapshot/SnapshotUpdaterImpl.java
+++ b/src/main/java/net/minestom/server/snapshot/SnapshotUpdaterImpl.java
@@ -8,7 +8,8 @@ import java.util.*;
 import java.util.concurrent.atomic.AtomicReference;
 
 final class SnapshotUpdaterImpl implements SnapshotUpdater {
-    private final Map<Snapshotable, AtomicReference<Snapshot>> referenceMap = new IdentityHashMap<>();
+    private final IdentityHashMap<Snapshotable, AtomicReference<Snapshot>> referenceMap = new IdentityHashMap<>();
+    private IdentityHashMap<Snapshotable, AtomicReference<Snapshot>> readOnlyReferenceMap;
     private List<Entry> queue = new ArrayList<>();
 
     static <T extends Snapshot> @NotNull T update(@NotNull Snapshotable snapshotable) {
@@ -19,12 +20,21 @@ final class SnapshotUpdaterImpl implements SnapshotUpdater {
     }
 
     @Override
-    public synchronized <T extends Snapshot> @NotNull AtomicReference<T> reference(@NotNull Snapshotable snapshotable) {
-        AtomicReference<Snapshot> ref = new AtomicReference<>();
-        var prev = referenceMap.putIfAbsent(snapshotable, ref);
-        if (prev != null) return (AtomicReference<T>) prev;
-        this.queue.add(new Entry(snapshotable, ref));
-        return (AtomicReference<T>) ref;
+    public <T extends Snapshot> @NotNull AtomicReference<T> reference(@NotNull Snapshotable snapshotable) {
+        AtomicReference<Snapshot> ref;
+        // Very often the same snapshotable is referenced multiple times.
+        var readOnly = this.readOnlyReferenceMap;
+        if (readOnly != null && (ref = readOnly.get(snapshotable)) != null) {
+            return (AtomicReference<T>) ref;
+        }
+        // If this is a new snapshotable, we need to create a new reference.
+        synchronized (this) {
+            ref = new AtomicReference<>();
+            var prev = referenceMap.putIfAbsent(snapshotable, ref);
+            if (prev != null) return (AtomicReference<T>) prev;
+            this.queue.add(new Entry(snapshotable, ref));
+            return (AtomicReference<T>) ref;
+        }
     }
 
     record Entry(Snapshotable snapshotable, AtomicReference<Snapshot> ref) {
@@ -34,6 +44,7 @@ final class SnapshotUpdaterImpl implements SnapshotUpdater {
         List<Entry> temp;
         while (!(temp = new ArrayList<>(queue)).isEmpty()) {
             queue = new ArrayList<>();
+            readOnlyReferenceMap = (IdentityHashMap<Snapshotable, AtomicReference<Snapshot>>) referenceMap.clone();
             temp.parallelStream().forEach(entry -> {
                 Snapshotable snap = entry.snapshotable;
                 entry.ref.set(Objects.requireNonNull(snap.updateSnapshot(this), "Snapshot must not be null after an update!"));