From 73dee1fb1a46e35b5863fc1027350fea2835418a Mon Sep 17 00:00:00 2001
From: TheMode <themode@outlook.fr>
Date: Wed, 21 Apr 2021 11:27:49 +0200
Subject: [PATCH] Better abstraction

---
 .../net/minestom/server/lock/Acquisition.java | 126 +++++++++++-------
 .../server/thread/ThreadProvider.java         |   1 -
 2 files changed, 80 insertions(+), 47 deletions(-)

diff --git a/src/main/java/net/minestom/server/lock/Acquisition.java b/src/main/java/net/minestom/server/lock/Acquisition.java
index a266bc3b4..0e0ea61ad 100644
--- a/src/main/java/net/minestom/server/lock/Acquisition.java
+++ b/src/main/java/net/minestom/server/lock/Acquisition.java
@@ -33,7 +33,7 @@ public final class Acquisition {
     public static <E> void acquireForEach(@NotNull Collection<Acquirable<E>> collection,
                                           @NotNull Consumer<? super E> consumer) {
         final Thread currentThread = Thread.currentThread();
-        Map<BatchThread, List<E>> threadCacheMap = retrieveThreadMap(collection, currentThread, consumer);
+        Map<BatchThread, List<E>> threadCacheMap = retrieveOptionalThreadMap(collection, currentThread, consumer);
 
         // Acquire all the threads one by one
         {
@@ -82,50 +82,61 @@ public final class Acquisition {
         if (Objects.equals(currentThread, elementThread)) {
             callback.run();
         } else {
-
-            // Monitoring
-            final boolean monitoring = MinecraftServer.hasWaitMonitoring();
-            long time = 0;
-            if (monitoring) {
-                time = System.nanoTime();
-            }
-
-            ReentrantLock currentLock;
-            {
-                final BatchThread current = currentThread instanceof BatchThread ?
-                        (BatchThread) currentThread : null;
-                currentLock = current != null && current.getLock().isHeldByCurrentThread() ?
-                        current.getLock() : null;
-            }
-            if (currentLock != null)
-                currentLock.unlock();
-
-            GLOBAL_LOCK.lock();
-
-            if (currentLock != null)
-                currentLock.lock();
-
-            final var lock = elementThread != null ? elementThread.getLock() : null;
-            final boolean acquired = lock == null || lock.isHeldByCurrentThread();
-            if (!acquired) {
-                lock.lock();
-            }
-
-            // Monitoring
-            if (monitoring) {
-                time = System.nanoTime() - time;
-                WAIT_COUNTER_NANO.addAndGet(time);
-            }
-
+            var lock = acquireEnter(currentThread, elementThread);
             callback.run();
-
-            if (!acquired) {
-                lock.unlock();
-            }
-            GLOBAL_LOCK.unlock();
+            acquireLeave(lock);
         }
     }
 
+    protected static ReentrantLock acquireEnter(Thread currentThread, BatchThread elementThread) {
+        // Monitoring
+        final boolean monitoring = MinecraftServer.hasWaitMonitoring();
+        long time = 0;
+        if (monitoring) {
+            time = System.nanoTime();
+        }
+
+        ReentrantLock currentLock;
+        {
+            final BatchThread current = currentThread instanceof BatchThread ?
+                    (BatchThread) currentThread : null;
+            currentLock = current != null && current.getLock().isHeldByCurrentThread() ?
+                    current.getLock() : null;
+        }
+        if (currentLock != null)
+            currentLock.unlock();
+
+        GLOBAL_LOCK.lock();
+
+        if (currentLock != null)
+            currentLock.lock();
+
+        final var lock = elementThread != null ? elementThread.getLock() : null;
+        final boolean acquired = lock == null || lock.isHeldByCurrentThread();
+        if (!acquired) {
+            lock.lock();
+        }
+
+        // Monitoring
+        if (monitoring) {
+            time = System.nanoTime() - time;
+            WAIT_COUNTER_NANO.addAndGet(time);
+        }
+
+        return !acquired ? lock : null;
+    }
+
+    protected static ReentrantLock acquireEnter(BatchThread elementThread) {
+        return acquireEnter(Thread.currentThread(), elementThread);
+    }
+
+    protected static void acquireLeave(ReentrantLock lock) {
+        if (lock != null) {
+            lock.unlock();
+        }
+        GLOBAL_LOCK.unlock();
+    }
+
     protected synchronized static <T> void scheduledAcquireRequest(@NotNull Acquirable<T> acquirable, Consumer<T> consumer) {
         ScheduledAcquisition scheduledAcquisition = SCHEDULED_ACQUISITION.get();
         scheduledAcquisition.acquirableElements.add((Acquirable<Object>) acquirable);
@@ -134,15 +145,23 @@ public final class Acquisition {
                 .add((Consumer<Object>) consumer);
     }
 
-    private static <E, T extends Acquirable<E>> Map<BatchThread, List<E>> retrieveThreadMap(@NotNull Collection<? super T> collection,
-                                                                                            @NotNull Thread currentThread,
-                                                                                            @NotNull Consumer<? super E> consumer) {
+    /**
+     * Creates
+     *
+     * @param collection    the acquirable collection
+     * @param currentThread the current thread
+     * @param consumer      the consumer to execute when an element is already in the current thread
+     * @param <E>           the acquirable element type
+     * @return a new Thread to acquirable elements map
+     */
+    protected static <E> Map<BatchThread, List<E>> retrieveOptionalThreadMap(@NotNull Collection<Acquirable<E>> collection,
+                                                                             @NotNull Thread currentThread,
+                                                                             @NotNull Consumer<? super E> consumer) {
         // Separate a collection of acquirable elements into a map of thread->elements
         // Useful to reduce the number of acquisition
 
         Map<BatchThread, List<E>> threadCacheMap = new HashMap<>();
-        for (Object obj : collection) {
-            T element = (T) obj;
+        for (Acquirable<E> element : collection) {
             final E value = element.unwrap();
 
             final BatchThread elementThread = element.getHandler().getBatchThread();
@@ -159,6 +178,21 @@ public final class Acquisition {
         return threadCacheMap;
     }
 
+    protected static <E> Map<BatchThread, List<E>> retrieveThreadMap(@NotNull Collection<Acquirable<E>> collection) {
+        // Separate a collection of acquirable elements into a map of thread->elements
+        // Useful to reduce the number of acquisition
+        Map<BatchThread, List<E>> threadCacheMap = new HashMap<>();
+        for (Acquirable<E> element : collection) {
+            final E value = element.unwrap();
+            final BatchThread elementThread = element.getHandler().getBatchThread();
+
+            List<E> threadCacheList = threadCacheMap.computeIfAbsent(elementThread, batchThread -> new ArrayList<>());
+            threadCacheList.add(value);
+        }
+
+        return threadCacheMap;
+    }
+
     public static long getCurrentWaitMonitoring() {
         return WAIT_COUNTER_NANO.get();
     }
diff --git a/src/main/java/net/minestom/server/thread/ThreadProvider.java b/src/main/java/net/minestom/server/thread/ThreadProvider.java
index fbb7dbe58..c64bb374d 100644
--- a/src/main/java/net/minestom/server/thread/ThreadProvider.java
+++ b/src/main/java/net/minestom/server/thread/ThreadProvider.java
@@ -221,7 +221,6 @@ public abstract class ThreadProvider {
                 break;
 
         }
-        System.out.println("update " + counter);
     }
 
     public void removeEntity(@NotNull Entity entity) {