Update & test acquirable api

Signed-off-by: TheMode <themode@outlook.fr>

src/main/java/net/minestom/server/ServerProcessImpl.java

@@ -1,6 +1,5 @@
 package net.minestom.server;
 
-import net.minestom.server.acquirable.Acquirable;
 import net.minestom.server.advancements.AdvancementManager;
 import net.minestom.server.adventure.bossbar.BossBarManager;
 import net.minestom.server.command.CommandManager;
@@ -27,6 +26,7 @@ import net.minestom.server.scoreboard.TeamManager;
 import net.minestom.server.storage.StorageLocation;
 import net.minestom.server.storage.StorageManager;
 import net.minestom.server.terminal.MinestomTerminal;
+import net.minestom.server.thread.Acquirable;
 import net.minestom.server.thread.ThreadDispatcher;
 import net.minestom.server.timer.SchedulerManager;
 import net.minestom.server.utils.PacketUtils;

src/main/java/net/minestom/server/entity/Entity.java

@@ -8,7 +8,6 @@ import net.kyori.adventure.text.event.HoverEventSource;
 import net.minestom.server.MinecraftServer;
 import net.minestom.server.Tickable;
 import net.minestom.server.Viewable;
-import net.minestom.server.acquirable.Acquirable;
 import net.minestom.server.collision.BoundingBox;
 import net.minestom.server.collision.CollisionUtils;
 import net.minestom.server.coordinate.Point;
@@ -37,6 +36,7 @@ import net.minestom.server.potion.PotionEffect;
 import net.minestom.server.potion.TimedPotion;
 import net.minestom.server.tag.Tag;
 import net.minestom.server.tag.TagHandler;
+import net.minestom.server.thread.Acquirable;
 import net.minestom.server.timer.Schedulable;
 import net.minestom.server.timer.Scheduler;
 import net.minestom.server.timer.TaskSchedule;
@@ -1530,11 +1530,6 @@ public class Entity implements Viewable, Tickable, Schedulable, TagHandler, Perm
         return (Acquirable<T>) acquirable;
     }
 
-    @ApiStatus.Experimental
-    public <T extends Entity> @NotNull Acquirable<T> getAcquirable(@NotNull Class<T> clazz) {
-        return (Acquirable<T>) acquirable;
-    }
-
     @Override
     public <T> @Nullable T getTag(@NotNull Tag<T> tag) {
         return tag.read(nbtCompound);

src/main/java/net/minestom/server/thread/Acquirable.java

@@ -1,8 +1,6 @@
-package net.minestom.server.acquirable;
+package net.minestom.server.thread;
 
 import net.minestom.server.entity.Entity;
-import net.minestom.server.thread.ThreadDispatcher;
-import net.minestom.server.thread.TickThread;
 import net.minestom.server.utils.async.AsyncUtils;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
@@ -12,7 +10,7 @@ import java.util.function.Consumer;
 import java.util.stream.Stream;
 
 @ApiStatus.Experimental
-public interface Acquirable<T> {
+public sealed interface Acquirable<T> permits AcquirableImpl {
 
     /**
      * Gets all the {@link Entity entities} being ticked in the current thread.
@@ -78,7 +76,7 @@ public interface Acquirable<T> {
      * @see #sync(Consumer) for auto-closeable capability
      */
     default @NotNull Acquired<T> lock() {
-        return new Acquired<>(unwrap(), getHandler().getTickThread());
+        return new Acquired<>(unwrap(), assignedThread());
     }
 
     /**
@@ -102,7 +100,7 @@ public interface Acquirable<T> {
      * @return true if the element is linked to the current thread
      */
     default boolean isLocal() {
-        return Thread.currentThread() == getHandler().getTickThread();
+        return Thread.currentThread() == assignedThread();
     }
 
     /**
@@ -139,32 +137,6 @@ public interface Acquirable<T> {
      */
     @NotNull T unwrap();
 
-    /**
-     * Gets the {@link Handler} of this acquirable element,
-     * containing the currently linked thread.
-     * <p>
-     * Mostly for internal use.
-     *
-     * @return this element handler
-     */
     @ApiStatus.Internal
-    @NotNull Handler getHandler();
-
-    final class Handler {
-        private volatile ThreadDispatcher.Partition partition;
-
-        public ThreadDispatcher.Partition getChunkEntry() {
-            return partition;
-        }
-
-        @ApiStatus.Internal
-        public void refreshChunkEntry(@NotNull ThreadDispatcher.Partition partition) {
-            this.partition = partition;
-        }
-
-        public TickThread getTickThread() {
-            final ThreadDispatcher.Partition entry = this.partition;
-            return entry != null ? entry.thread() : null;
-        }
-    }
+    @NotNull TickThread assignedThread();
 }

src/main/java/net/minestom/server/thread/AcquirableCollection.java

@@ -1,6 +1,5 @@
-package net.minestom.server.acquirable;
+package net.minestom.server.thread;
 
-import net.minestom.server.thread.TickThread;
 import net.minestom.server.utils.async.AsyncUtils;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
@@ -11,7 +10,6 @@ import java.util.stream.Stream;
 
 @ApiStatus.Experimental
 public class AcquirableCollection<E> implements Collection<Acquirable<E>> {
-
     private final Collection<Acquirable<E>> acquirableCollection;
 
     public AcquirableCollection(Collection<Acquirable<E>> acquirableCollection) {
@@ -131,7 +129,7 @@ public class AcquirableCollection<E> implements Collection<Acquirable<E>> {
         for (var element : collection) {
             final T value = element.unwrap();
 
-            final TickThread elementThread = element.getHandler().getTickThread();
+            final TickThread elementThread = element.assignedThread();
             if (currentThread == elementThread) {
                 // The element is managed in the current thread, consumer can be immediately called
                 consumer.accept(value);

src/main/java/net/minestom/server/thread/AcquirableImpl.java

@@ -1,14 +1,24 @@
-package net.minestom.server.acquirable;
+package net.minestom.server.thread;
 
-import net.minestom.server.thread.TickThread;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.ReentrantLock;
 
 final class AcquirableImpl<T> implements Acquirable<T> {
     static final AtomicLong WAIT_COUNTER_NANO = new AtomicLong();
+    private static final VarHandle ASSIGNED_THREAD;
+
+    static {
+        try {
+            ASSIGNED_THREAD = MethodHandles.lookup().findVarHandle(AcquirableImpl.class, "assignedThread", TickThread.class);
+        } catch (NoSuchFieldException | IllegalAccessException e) {
+            throw new IllegalStateException(e);
+        }
+    }
 
     /**
      * Global lock used for synchronization.
@@ -16,11 +26,11 @@ final class AcquirableImpl<T> implements Acquirable<T> {
     private static final ReentrantLock GLOBAL_LOCK = new ReentrantLock();
 
     private final T value;
-    private final Acquirable.Handler handler;
+    @SuppressWarnings("unused")
+    private TickThread assignedThread;
 
     public AcquirableImpl(@NotNull T value) {
         this.value = value;
-        this.handler = new Acquirable.Handler();
     }
 
     @Override
@@ -29,8 +39,12 @@ final class AcquirableImpl<T> implements Acquirable<T> {
     }
 
     @Override
-    public @NotNull Acquirable.Handler getHandler() {
-        return handler;
+    public @NotNull TickThread assignedThread() {
+        return (TickThread) ASSIGNED_THREAD.getAcquire(this);
+    }
+
+    void updateThread(@NotNull TickThread thread) {
+        ASSIGNED_THREAD.setRelease(this, thread);
     }
 
     static @Nullable ReentrantLock enter(@NotNull Thread currentThread, @Nullable TickThread elementThread) {

src/main/java/net/minestom/server/thread/Acquired.java

@@ -1,6 +1,5 @@
-package net.minestom.server.acquirable;
+package net.minestom.server.thread;
 
-import net.minestom.server.thread.TickThread;
 import net.minestom.server.utils.validate.Check;
 import org.jetbrains.annotations.NotNull;
 

src/main/java/net/minestom/server/thread/ThreadDispatcher.java

@@ -1,7 +1,7 @@
 package net.minestom.server.thread;
 
 import net.minestom.server.Tickable;
-import net.minestom.server.acquirable.Acquirable;
+import net.minestom.server.entity.Entity;
 import org.jctools.queues.MessagePassingQueue;
 import org.jctools.queues.MpscUnboundedArrayQueue;
 import org.jetbrains.annotations.ApiStatus;
@@ -196,8 +196,8 @@ public final class ThreadDispatcher<P> {
         if (partitionEntry != null) {
             this.elements.put(tickable, partitionEntry);
             partitionEntry.elements.add(tickable);
-            if (tickable instanceof Acquirable<?> acquirable) {
-                acquirable.getHandler().refreshChunkEntry(partitionEntry);
+            if (tickable instanceof Entity entity) { // TODO support other types
+                ((AcquirableImpl<?>) entity.getAcquirable()).updateThread(partitionEntry.thread());
             }
         }
     }

src/test/java/net/minestom/server/thread/AcquirableTest.java

@@ -0,0 +1,44 @@
+package net.minestom.server.thread;
+
+import net.minestom.server.entity.Entity;
+import net.minestom.server.entity.EntityType;
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+public class AcquirableTest {
+
+    @Test
+    public void assignation() {
+        AtomicReference<TickThread> tickThread = new AtomicReference<>();
+        Entity entity = new Entity(EntityType.ZOMBIE) {
+            @Override
+            public void tick(long time) {
+                super.tick(time);
+                tickThread.set(getAcquirable().assignedThread());
+            }
+        };
+        Object first = new Object();
+        Object second = new Object();
+
+        ThreadDispatcher<Object> dispatcher = ThreadDispatcher.of(ThreadProvider.counter(), 2);
+        dispatcher.createPartition(first);
+        dispatcher.createPartition(second);
+
+        dispatcher.updateElement(entity, first);
+        dispatcher.updateAndAwait(System.currentTimeMillis());
+        TickThread firstThread = tickThread.get();
+        assertNotNull(firstThread);
+
+        tickThread.set(null);
+        dispatcher.updateElement(entity, second);
+        dispatcher.updateAndAwait(System.currentTimeMillis());
+        TickThread secondThread = tickThread.get();
+        assertNotNull(secondThread);
+
+        assertNotEquals(firstThread, secondThread);
+    }
+}