Implement gradual chunk loading after login

gradle.properties

@@ -2,7 +2,7 @@
 modId=viewdistanceworkaround
 modName=View Distance Workaround
 modAuthor=chylex
-modVersion=1.0.0
+modVersion=1.1.0
 
 # Dependencies
 minecraftVersion=1.17.1

src/main/java/chylex/viewdistance/mixin/ChunkMapAccessor.java

@@ -0,0 +1,14 @@
+package chylex.viewdistance.mixin;
+
+import net.minecraft.network.protocol.Packet;
+import net.minecraft.server.level.ChunkMap;
+import net.minecraft.server.level.ServerPlayer;
+import net.minecraft.world.level.ChunkPos;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Invoker;
+
+@Mixin(ChunkMap.class)
+public interface ChunkMapAccessor {
+	@Invoker
+	void callUpdateChunkTracking(final ServerPlayer serverPlayer, final ChunkPos pos, final Packet<?>[] packetArrayForReuse, final boolean stopTracking, final boolean startTracking);
+}

src/main/java/chylex/viewdistance/mixin/ChunkMapMixin.java

@@ -1,5 +1,6 @@
 package chylex.viewdistance.mixin;
 import chylex.viewdistance.ViewDistanceWorkaround;
+import chylex.viewdistance.track.ChunkTrackerThread;
 import net.minecraft.server.dedicated.DedicatedServer;
 import net.minecraft.server.level.ChunkMap;
 import net.minecraft.server.level.ServerPlayer;
@@ -9,28 +10,31 @@ import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.Redirect;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
-import java.util.HashSet;
-import java.util.Set;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.UUID;
 
 @Mixin(ChunkMap.class)
 public class ChunkMapMixin {
-	private final Set<UUID> players = new HashSet<>();
+	private final Map<UUID, ChunkTrackerThread> playerThreads = new HashMap<>();
 	private boolean active = false;
 	
 	@Shadow
 	int viewDistance;
 	
-	@Shadow
-	void updatePlayerStatus(final ServerPlayer serverPlayer, final boolean isRemoving) {}
-	
 	@Inject(method = "updatePlayerStatus", at = @At("HEAD"))
 	private void hookUpdatePlayerStatus(final ServerPlayer serverPlayer, final boolean isAdding, final CallbackInfo ci) {
+		final UUID uuid = serverPlayer.getUUID();
+		
 		if (!isAdding) {
-			players.remove(serverPlayer.getUUID());
+			var thread = playerThreads.remove(uuid);
+			if (thread != null) {
+				thread.interrupt();
+			}
+			
 			active = false;
 		}
-		else if (players.add(serverPlayer.getUUID())) {
+		else if (!playerThreads.containsKey(uuid)) {
 			final DedicatedServer dedi = ViewDistanceWorkaround.get().getDedicatedServer();
 			if (dedi == null) {
 				active = false;
@@ -39,27 +43,24 @@ public class ChunkMapMixin {
 			
 			active = true;
 			
-			new Thread(() -> {
-				try {
-					
-					Thread.sleep(2_000L);
-					dedi.execute(() -> {
-						if (players.contains(serverPlayer.getUUID())) {
-							updatePlayerStatus(serverPlayer, isAdding);
-						}
-					});
-					
-				} catch (final InterruptedException ignored) {
-				}
-			}).start();
+			@SuppressWarnings("ConstantConditions")
+			final ChunkMap chunkMap = (ChunkMap)(Object)this;
+			final ChunkTrackerThread thread = new ChunkTrackerThread(dedi, chunkMap, serverPlayer, viewDistance);
+			playerThreads.put(uuid, thread);
+			thread.start();
 		}
 		else {
 			active = false;
 		}
 	}
 	
+	@Inject(method = "updatePlayerStatus", at = @At("RETURN"))
+	private void hookUpdatePlayerStatusEnd(final ServerPlayer serverPlayer, final boolean isAdding, final CallbackInfo ci) {
+		active = false;
+	}
+	
 	@Redirect(method = "updatePlayerStatus", at = @At(value = "FIELD", target = "Lnet/minecraft/server/level/ChunkMap;viewDistance:I"))
 	private int getViewDistance(final ChunkMap instance) {
-		return active ? 2 : viewDistance;
+		return active ? 1 : viewDistance;
 	}
 }

src/main/java/chylex/viewdistance/track/ChunkTrackerThread.java

@@ -0,0 +1,74 @@
+package chylex.viewdistance.track;
+
+import chylex.viewdistance.mixin.ChunkMapAccessor;
+import net.minecraft.core.SectionPos;
+import net.minecraft.server.dedicated.DedicatedServer;
+import net.minecraft.server.level.ChunkMap;
+import net.minecraft.server.level.ServerPlayer;
+import net.minecraft.world.level.ChunkPos;
+import org.apache.commons.lang3.mutable.MutableObject;
+import java.util.PriorityQueue;
+import java.util.Queue;
+import static java.util.Comparator.comparingInt;
+
+public class ChunkTrackerThread extends Thread {
+	private final DedicatedServer dedi;
+	private final ChunkMapAccessor chunkMap;
+	private final ServerPlayer player;
+	private final Queue<ChunkPos> remainingChunks;
+	
+	public ChunkTrackerThread(final DedicatedServer dedi, final ChunkMap chunkMap, final ServerPlayer player, final int viewDistance) {
+		super("ViewDistanceWorkaround / " + player.getScoreboardName());
+		
+		this.dedi = dedi;
+		this.chunkMap = (ChunkMapAccessor) chunkMap;
+		this.player = player;
+		
+		final int centerChunkX = SectionPos.blockToSectionCoord(player.getBlockX());
+		final int centerChunkZ = SectionPos.blockToSectionCoord(player.getBlockZ());
+		final ChunkPos centerChunk = new ChunkPos(centerChunkX, centerChunkZ);
+		
+		this.remainingChunks = new PriorityQueue<>(comparingInt(pos -> distanceSquared(pos, centerChunk)));
+		
+		for(int x = centerChunkX - viewDistance; x <= centerChunkX + viewDistance; ++x) {
+			for(int z = centerChunkZ - viewDistance; z <= centerChunkZ + viewDistance; ++z) {
+				this.remainingChunks.add(new ChunkPos(x, z));
+			}
+		}
+	}
+	
+	@SuppressWarnings("BusyWait")
+	@Override
+	public void run() {
+		try {
+			Thread.sleep(1250L);
+		} catch (InterruptedException e) {
+			return;
+		}
+		
+		System.out.println("[ViewDistanceWorkaround] Starting delayed chunk tracking for " + player.getScoreboardName() + " (" + remainingChunks.size() + " chunks)");
+		
+		ChunkPos nextChunk;
+		while ((nextChunk = remainingChunks.poll()) != null && player.connection.connection.isConnected()) {
+			ChunkPos pos = nextChunk;
+			
+			dedi.executeBlocking(() -> {
+				chunkMap.callUpdateChunkTracking(player, pos, new MutableObject<>(), false, true);
+			});
+			
+			try {
+				Thread.sleep(5L);
+			} catch (InterruptedException e) {
+				return;
+			}
+		}
+		
+		System.out.println("[ViewDistanceWorkaround] Finished delayed chunk tracking for " + player.getScoreboardName());
+	}
+	
+	private static int distanceSquared(ChunkPos p1, ChunkPos p2) {
+		int xDiff = p1.x - p2.x;
+		int zDiff = p1.z - p2.z;
+		return (xDiff * xDiff) + (zDiff * zDiff);
+	}
+}

src/main/resources/viewdistanceworkaround.mixins.json

@@ -5,6 +5,7 @@
   "refmap": "viewdistanceworkaround.refmap.json",
   "compatibilityLevel": "JAVA_16",
   "server": [
+    "ChunkMapAccessor",
     "ChunkMapMixin",
     "CommandMixin"
   ],