From 3211831073c350a087a13b542c4c9919c571cbf3 Mon Sep 17 00:00:00 2001
From: TheMode <themode@outlook.fr>
Date: Thu, 16 Sep 2021 07:43:08 +0200
Subject: [PATCH] Ensure that the viewable element can be collected

Signed-off-by: TheMode <themode@outlook.fr>
---
 .../minestom/server/utils/PacketUtils.java    | 36 ++++++++++++-------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/src/main/java/net/minestom/server/utils/PacketUtils.java b/src/main/java/net/minestom/server/utils/PacketUtils.java
index 96e9d0366..ec1726134 100644
--- a/src/main/java/net/minestom/server/utils/PacketUtils.java
+++ b/src/main/java/net/minestom/server/utils/PacketUtils.java
@@ -23,6 +23,7 @@ import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.lang.ref.WeakReference;
 import java.nio.ByteBuffer;
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
@@ -261,12 +262,12 @@ public final class PacketUtils {
     }
 
     private static final class ViewableStorage {
-        private final Viewable viewable;
+        private final WeakReference<Viewable> viewable;
         private final Map<PlayerConnection, List<IntIntPair>> entityIdMap = new HashMap<>();
         private final BinaryBuffer buffer = PooledBuffers.get();
 
         private ViewableStorage(Viewable viewable) {
-            this.viewable = viewable;
+            this.viewable = new WeakReference<>(viewable);
             PooledBuffers.registerBuffer(this, buffer);
         }
 
@@ -274,7 +275,7 @@ public final class PacketUtils {
             final ByteBuffer framedPacket = createFramedPacket(serverPacket).flip();
             final int packetSize = framedPacket.limit();
             if (packetSize >= buffer.capacity()) {
-                processSingle(framedPacket, connection);
+                process(new SingleEntry(framedPacket, connection));
                 return;
             }
             if (!buffer.canWrite(packetSize)) process();
@@ -287,9 +288,11 @@ public final class PacketUtils {
             }
         }
 
-        private synchronized void process() {
+        private synchronized void process(@Nullable SingleEntry singleEntry) {
             if (buffer.writerOffset() == 0)
                 return; // TODO: there is nothing in the buffer, remove from VIEWABLE_STORAGE_MAP
+            final Viewable viewable = this.viewable.get();
+            if (viewable == null) return;
             for (Player player : viewable.getViewers()) {
                 PlayerConnection connection = player.getPlayerConnection();
                 Consumer<ByteBuffer> writer = connection instanceof PlayerSocketConnection
@@ -316,21 +319,28 @@ public final class PacketUtils {
                     ByteBuffer remainSlice = buffer.asByteBuffer(lastWrite, remaining);
                     writer.accept(remainSlice);
                 }
+
+                // Handle single entry
+                if (singleEntry != null && !singleEntry.exception.equals(connection)) {
+                    writer.accept(singleEntry.buffer.position(0));
+                }
             }
             // Clear state
             this.entityIdMap.clear();
             this.buffer.clear();
         }
 
-        private synchronized void processSingle(ByteBuffer buffer, PlayerConnection exception) {
-            process();
-            for (Player player : viewable.getViewers()) {
-                PlayerConnection connection = player.getPlayerConnection();
-                if (Objects.equals(connection, exception)) continue;
-                if (connection instanceof PlayerSocketConnection) {
-                    ((PlayerSocketConnection) connection).write(buffer.position(0));
-                }
-                // TODO for non-socket connection
+        private void process() {
+            process(null);
+        }
+
+        private static final class SingleEntry {
+            private final ByteBuffer buffer;
+            private final PlayerConnection exception;
+
+            public SingleEntry(ByteBuffer buffer, PlayerConnection exception) {
+                this.buffer = buffer;
+                this.exception = exception;
             }
         }
     }