Add a config to specify chunk recorder radius (#28)

build.gradle.kts

@@ -28,6 +28,8 @@ repositories {
 }
 
 val minecraftVersion: String by project
+// Different name for loom...
+val mcVersion = minecraftVersion
 val parchmentVersion: String by project
 val loaderVersion: String by project
 val fabricVersion: String by project
@@ -81,7 +83,7 @@ loom {
 
     runs {
         getByName("server") {
-            runDir = "run/${minecraftVersion}"
+            runDir = "run/$mcVersion"
         }
 
         getByName("client") {

src/main/java/me/senseiwells/replay/mixin/chunk/ChunkHolderMixin.java

@@ -1,15 +1,15 @@
 package me.senseiwells.replay.mixin.chunk;
 
 import com.llamalad7.mixinextras.injector.ModifyExpressionValue;
+import com.mojang.datafixers.util.Either;
 import me.senseiwells.replay.chunk.ChunkRecorder;
 import me.senseiwells.replay.ducks.ServerReplay$ChunkRecordable;
 import net.minecraft.network.protocol.Packet;
 import net.minecraft.server.level.ChunkHolder;
+import net.minecraft.server.level.ChunkMap;
 import net.minecraft.server.level.ServerPlayer;
-import net.minecraft.world.level.ChunkPos;
 import net.minecraft.world.level.chunk.LevelChunk;
 import org.jetbrains.annotations.Nullable;
-import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.Unique;
@@ -21,12 +21,16 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
 
 @Mixin(ChunkHolder.class)
 public abstract class ChunkHolderMixin implements ServerReplay$ChunkRecordable {
+	@Shadow private volatile CompletableFuture<Either<LevelChunk, ChunkHolder.ChunkLoadingFailure>> fullChunkFuture;
+
 	@Unique private final Set<ChunkRecorder> replay$recorders = new HashSet<>();
 
-	@Shadow @Final ChunkPos pos;
+	@Shadow public abstract @Nullable LevelChunk getFullChunk();
 
 	@Inject(
 		method = "broadcast",
@@ -50,6 +54,42 @@ private boolean shouldSkipBroadcasting(boolean noPlayers) {
 		return noPlayers && this.replay$recorders.isEmpty();
 	}
 
+	@Inject(
+		method = "updateFutures",
+		at = @At(
+			value = "INVOKE",
+			target = "Lnet/minecraft/server/level/ChunkHolder;updateChunkToSave(Ljava/util/concurrent/CompletableFuture;Ljava/lang/String;)V",
+			shift = At.Shift.AFTER,
+			ordinal = 0
+		)
+	)
+	private void onChunkLoad(ChunkMap chunkMap, Executor executor, CallbackInfo ci) {
+		this.fullChunkFuture.thenAccept(result -> {
+			result.ifLeft(chunk -> {
+				for (ChunkRecorder recorder : this.getRecorders()) {
+					recorder.onChunkLoaded(chunk);
+				}
+			});
+		});
+	}
+
+	@Inject(
+		method = "updateFutures",
+		at = @At(
+			value = "INVOKE",
+			target = "Ljava/util/concurrent/CompletableFuture;complete(Ljava/lang/Object;)Z",
+			ordinal = 0
+		)
+	)
+	private void onChunkUnload(ChunkMap chunkMap, Executor executor, CallbackInfo ci) {
+		LevelChunk chunk = this.getFullChunk();
+		if (chunk != null) {
+			for (ChunkRecorder recorder : this.getRecorders()) {
+				recorder.onChunkUnloaded(chunk);
+			}
+		}
+	}
+
 	@Override
 	public Collection<ChunkRecorder> replay$getRecorders() {
 		return this.replay$recorders;
@@ -58,23 +98,33 @@ private boolean shouldSkipBroadcasting(boolean noPlayers) {
 	@Override
 	public void replay$addRecorder(ChunkRecorder recorder) {
 		if (this.replay$recorders.add(recorder)) {
-			recorder.onChunkLoaded(this.pos);
+			this.fullChunkFuture.thenAccept(result -> {
+				result.ifLeft(recorder::onChunkLoaded);
+			});
+
 			recorder.addRecordable(this);
 		}
 	}
 
 	@Override
 	public void replay$removeRecorder(ChunkRecorder recorder) {
 		if (this.replay$recorders.remove(recorder)) {
-			recorder.onChunkUnloaded(this.pos);
+			LevelChunk chunk = this.getFullChunk();
+			if (chunk != null) {
+				recorder.onChunkUnloaded(chunk);
+			}
+
 			recorder.removeRecordable(this);
 		}
 	}
 
 	@Override
 	public void replay$removeAllRecorders() {
+		LevelChunk chunk = this.getFullChunk();
 		for (ChunkRecorder recorder : this.replay$recorders) {
-			recorder.onChunkUnloaded(this.pos);
+			if (chunk != null) {
+				recorder.onChunkUnloaded(chunk);
+			}
 			recorder.removeRecordable(this);
 		}
 		this.replay$recorders.clear();

src/main/kotlin/me/senseiwells/replay/chunk/ChunkRecorder.kt

@@ -1,7 +1,6 @@
 package me.senseiwells.replay.chunk
 
 import it.unimi.dsi.fastutil.ints.IntArraySet
-import it.unimi.dsi.fastutil.ints.IntSet
 import it.unimi.dsi.fastutil.longs.LongOpenHashSet
 import kotlinx.serialization.json.JsonElement
 import kotlinx.serialization.json.JsonPrimitive
@@ -25,6 +24,7 @@ import net.minecraft.server.level.ServerPlayer
 import net.minecraft.world.entity.Entity
 import net.minecraft.world.entity.boss.wither.WitherBoss
 import net.minecraft.world.level.ChunkPos
+import net.minecraft.world.level.chunk.LevelChunk
 import net.minecraft.world.level.levelgen.Heightmap
 import org.apache.commons.lang3.builder.ToStringBuilder
 import org.jetbrains.annotations.ApiStatus.Internal
@@ -48,10 +48,11 @@ class ChunkRecorder internal constructor(
     val recorderName: String,
     recordings: Path
 ): ReplayRecorder(chunks.level.server, PROFILE, recordings), ChunkSender {
-    private val dummy = ServerPlayer(this.server, this.chunks.level, PROFILE, ClientInformation.createDefault())
+    private val dummy by lazy {
+        ServerPlayer(this.server, this.chunks.level, PROFILE, ClientInformation.createDefault())
+    }
 
     private val sentChunks = LongOpenHashSet()
-    private var sentAllChunks = false
 
     private val recordables = HashSet<ChunkRecordable>()
 
@@ -95,13 +96,7 @@ class ChunkRecorder internal constructor(
 
         RejoinedReplayPlayer.rejoin(this.dummy, this)
         this.spawnPlayer()
-
-        if (ServerReplay.config.loadAllChunkRecorderChunks) {
-            this.sentAllChunks = true
-            this.sendChunksAndEntities()
-        } else {
-            this.sendChunkViewDistance()
-        }
+        this.sendChunksAndEntities()
 
         val chunks = this.level.chunkSource.chunkMap as ChunkMapAccessor
         for (pos in this.chunks) {
@@ -196,7 +191,15 @@ class ChunkRecorder internal constructor(
      * @param consumer The consumer that will accept the given chunks positions.
      */
     override fun forEachChunk(consumer: Consumer<ChunkPos>) {
-        this.chunks.forEach(consumer)
+        val radius = ServerReplay.config.chunkRecorderLoadRadius
+        if (radius < 0) {
+            this.chunks.forEach(consumer)
+            return
+        }
+
+        ChunkPos.rangeClosed(this.chunks.center, radius + 1).filter {
+            this.chunks.contains(this.level.dimension(), it)
+        }.forEach(consumer)
     }
 
     /**
@@ -238,6 +241,10 @@ class ChunkRecorder internal constructor(
         return this.chunks.viewDistance
     }
 
+    override fun onChunkSent(chunk: LevelChunk) {
+        this.sentChunks.add(chunk.pos.toLong())
+    }
+
     /**
      * Determines whether a given packet is able to be recorded.
      *
@@ -264,7 +271,7 @@ class ChunkRecorder internal constructor(
      *
      * @return The dummy chunk recording player.
      */
-    fun getDummy(): ServerPlayer {
+    fun getDummyPlayer(): ServerPlayer {
         return this.dummy
     }
 
@@ -295,24 +302,25 @@ class ChunkRecorder internal constructor(
     }
 
     @Internal
-    fun onChunkLoaded(pos: ChunkPos) {
-        if (!this.chunks.contains(this.level.dimension(), pos)) {
+    fun onChunkLoaded(chunk: LevelChunk): Boolean {
+        if (!this.chunks.contains(chunk.level.dimension(), chunk.pos)) {
             ServerReplay.logger.error("Tried to load chunk out of bounds!")
-            return
+            return false
         }
 
         this.loadedChunks++
         this.resume()
 
-        if (!this.sentAllChunks && this.sentChunks.add(pos.toLong())) {
-            val chunk = this.level.getChunk(pos.x, pos.z)
+        if (!this.sentChunks.contains(chunk.pos.toLong())) {
             this.sendChunk(this.level.chunkSource.chunkMap, chunk, IntArraySet())
         }
+
+        return true
     }
 
     @Internal
-    fun onChunkUnloaded(pos: ChunkPos) {
-        if (!this.chunks.contains(this.level.dimension(), pos)) {
+    fun onChunkUnloaded(chunk: LevelChunk) {
+        if (!this.chunks.contains(chunk.level.dimension(), chunk.pos)) {
             ServerReplay.logger.error("Tried to unload chunk out of bounds!")
             return
         }

src/main/kotlin/me/senseiwells/replay/compat/voicechat/ReplayVoicechatPlugin.kt

@@ -245,7 +245,7 @@ object ReplayVoicechatPlugin: VoicechatPlugin, ServerReplayPlugin {
         this.recordAdditionalPackets(recorder)
         val server = Voicechat.SERVER.server
         if (server != null) {
-            val player = recorder.getDummy()
+            val player = recorder.getDummyPlayer()
             // The chunks aren't sending any voice data so doesn't need a secret
             val packet = SecretPacket(player, Util.NIL_UUID, server.port, Voicechat.SERVER_CONFIG)
             recorder.record(packet.toClientboundPacket())

src/main/kotlin/me/senseiwells/replay/config/ReplayConfig.kt

@@ -68,8 +68,8 @@ class ReplayConfig {
     @SerialName("fixed_daylight_cycle")
     var fixedDaylightCycle = -1L
 
-    @SerialName("load_all_chunk_recorder_chunks")
-    var loadAllChunkRecorderChunks = true
+    @SerialName("chunk_recorder_load_radius")
+    var chunkRecorderLoadRadius = -1
     @SerialName("pause_unloaded_chunks")
     var skipWhenChunksUnloaded = false
     @SerialName("pause_notify_players")

src/main/kotlin/me/senseiwells/replay/recorder/ChunkSender.kt

@@ -19,6 +19,7 @@ import net.minecraft.world.level.ChunkPos
 import net.minecraft.world.level.chunk.LevelChunk
 import org.jetbrains.annotations.ApiStatus.Internal
 import org.jetbrains.annotations.ApiStatus.NonExtendable
+import org.jetbrains.annotations.ApiStatus.OverrideOnly
 import java.util.function.Consumer
 import kotlin.math.min
 
@@ -85,6 +86,16 @@ interface ChunkSender {
         return this.level.server.playerList.viewDistance
     }
 
+    /**
+     * This is called when a chunk is successfully sent to the client.
+     *
+     * @param chunk The chunk that was sent.
+     */
+    @OverrideOnly
+    fun onChunkSent(chunk: LevelChunk) {
+
+    }
+
     /**
      * This sends all chunk and entity packets.
      */
@@ -180,6 +191,8 @@ interface ChunkSender {
         for (entity in ridden) {
             this.sendPacket(ClientboundSetPassengersPacket(entity))
         }
+
+        this.onChunkSent(chunk)
     }
 
     /**

src/main/kotlin/me/senseiwells/replay/recorder/ReplayRecorder.kt

@@ -92,7 +92,8 @@ abstract class ReplayRecorder(
 
     private var packId = 0
 
-    private var started = false
+    internal var started = false
+        private set
 
     private var ignore = false
 

src/main/resources/serverreplay.mixins.json

@@ -18,6 +18,7 @@
     "common.ServerLevelMixin",
     "compat.carpet.EntityPlayerMPFakeMixin",
     "compat.carpet.NetHandlerPlayServerFakeMixin",
+    "compat.vmp.NearbyEntityTrackingMixin",
     "compat.vmp.TrackedEntityMixin",
     "player.ServerCommonPacketListenerImplMixin",
     "player.ServerConfigurationPacketListenerImplMixin",