Reimplement block entity (un)load events (#1177)

* Reimplement block entity (un)load events

This is pending testing to verify the tracking is reliable.

* Forgot that null check for parity

* Everything is implemented. Now for testing

* Fix server block entity unload events

* Bah indentations broke

* Handle world chunks instead of the positions

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/api/event/lifecycle/v1/ServerWorldEvents.java

@@ -25,9 +25,9 @@ import net.fabricmc.fabric.api.event.EventFactory;
 
 public final class ServerWorldEvents {
 	/**
-	 * Called when a world is loaded by a Minecraft server.
+	 * Called just after a world is loaded by a Minecraft server.
 	 *
-	 * <p>For example, this can be used to load world specific metadata or initialize a {@link PersistentState} on a server world.
+	 * <p>This can be used to load world specific metadata or initialize a {@link PersistentState} on a server world.
 	 */
 	public static final Event<Load> LOAD = EventFactory.createArrayBacked(Load.class, callbacks -> (server, world) -> {
 		for (Load callback : callbacks) {
@@ -39,7 +39,7 @@ public final class ServerWorldEvents {
 	 * Called before a world is unloaded by a Minecraft server.
 	 *
 	 * <p>This typically occurs after a server has {@link ServerLifecycleEvents#SERVER_STOPPING started shutting down}.
-	 * Mods which allow dynamic world (un)registration should use this event so mods can let go of world handles when a world is removed.
+	 * Mods which allow dynamic world (un)registration should call this event so mods can let go of world handles when a world is removed.
 	 */
 	public static final Event<Unload> UNLOAD = EventFactory.createArrayBacked(Unload.class, callbacks -> (server, world) -> {
 		for (Unload callback : callbacks) {

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/impl/event/lifecycle/ClientLifecycleEventsImpl.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.impl.event.lifecycle;
+
+import net.minecraft.block.entity.BlockEntity;
+
+import net.fabricmc.api.ClientModInitializer;
+import net.fabricmc.api.EnvType;
+import net.fabricmc.api.Environment;
+import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientBlockEntityEvents;
+import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientChunkEvents;
+
+@Environment(EnvType.CLIENT)
+public final class ClientLifecycleEventsImpl implements ClientModInitializer {
+	@Override
+	public void onInitializeClient() {
+		// Part of impl for block entity events
+		ClientChunkEvents.CHUNK_LOAD.register((world, chunk) -> {
+			((LoadedChunksCache) world).fabric_markLoaded(chunk);
+		});
+
+		ClientChunkEvents.CHUNK_UNLOAD.register((world, chunk) -> {
+			((LoadedChunksCache) world).fabric_markUnloaded(chunk);
+		});
+
+		ClientChunkEvents.CHUNK_UNLOAD.register((world, chunk) -> {
+			for (BlockEntity blockEntity : chunk.getBlockEntities().values()) {
+				ClientBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, world);
+			}
+		});
+	}
+}

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/impl/event/lifecycle/LifecycleEventsImpl.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.impl.event.lifecycle;
+
+import net.minecraft.block.entity.BlockEntity;
+import net.minecraft.world.chunk.WorldChunk;
+
+import net.fabricmc.api.ModInitializer;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerBlockEntityEvents;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerChunkEvents;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerWorldEvents;
+
+public final class LifecycleEventsImpl implements ModInitializer {
+	@Override
+	public void onInitialize() {
+		// Part of impl for block entity events
+		ServerChunkEvents.CHUNK_LOAD.register((world, chunk) -> {
+			((LoadedChunksCache) world).fabric_markLoaded(chunk);
+		});
+
+		ServerChunkEvents.CHUNK_UNLOAD.register((world, chunk) -> {
+			((LoadedChunksCache) world).fabric_markUnloaded(chunk);
+		});
+
+		// Fire block entity unload events.
+		// This handles the edge case where going through a portal will cause block entities to unload without warning.
+		ServerChunkEvents.CHUNK_UNLOAD.register((world, chunk) -> {
+			for (BlockEntity blockEntity : chunk.getBlockEntities().values()) {
+				ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, world);
+			}
+		});
+
+		// We use the world unload event so worlds that are dynamically hot(un)loaded get block entity unload events fired when shut down.
+		ServerWorldEvents.UNLOAD.register((server, world) -> {
+			for (WorldChunk chunk : ((LoadedChunksCache) world).fabric_getLoadedChunks()) {
+				for (BlockEntity blockEntity : chunk.getBlockEntities().values()) {
+					ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, world);
+				}
+			}
+		});
+	}
+}

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/impl/event/lifecycle/LoadedChunksCache.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.impl.event.lifecycle;
+
+import java.util.Set;
+
+import net.minecraft.world.chunk.WorldChunk;
+
+/**
+ * A simple marker interface which holds references to chunks which block entities may be loaded or unloaded from.
+ */
+public interface LoadedChunksCache {
+	Set<WorldChunk> fabric_getLoadedChunks();
+
+	/**
+	 * Marks a chunk as loaded in a world.
+	 */
+	void fabric_markLoaded(WorldChunk chunk);
+
+	/**
+	 * Marks a chunk as unloaded in a world.
+	 */
+	void fabric_markUnloaded(WorldChunk chunk);
+}

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/mixin/event/lifecycle/MinecraftServerMixin.java

@@ -74,16 +74,6 @@ public abstract class MinecraftServerMixin {
 		ServerTickEvents.END_SERVER_TICK.invoker().onEndTick((MinecraftServer) (Object) this);
 	}
 
-	/**
-	 * When a world is closed, it means the world will be unloaded.
-	 */
-	/*@Inject(at = @At(value = "INVOKE", target = "Lnet/minecraft/server/world/ServerWorld;close()V"), method = "shutdown", locals = LocalCapture.CAPTURE_FAILEXCEPTION)
-	private void closeWorld(CallbackInfo ci, Iterator<ServerWorld> worlds, ServerWorld serverWorld) {
-		for (BlockEntity blockEntity : serverWorld.blockEntities) {
-			ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, serverWorld);
-		}
-	}*/
-
 	// The locals you have to manage for an inject are insane. And do it twice. A redirect is much cleaner.
 	// Here is what it looks like with an inject: https://gist.github.com/i509VCB/f80077cc536eb4dba62b794eba5611c1
 	@Redirect(method = "createWorlds", at = @At(value = "INVOKE", target = "Ljava/util/Map;put(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;"))
@@ -94,7 +84,7 @@ public abstract class MinecraftServerMixin {
 		return result;
 	}
 
-	@Inject(method = "shutdown", at = @At(value = "INVOKE", target = "Lnet/minecraft/server/world/ServerWorld;close()V", shift = At.Shift.AFTER), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+	@Inject(method = "shutdown", at = @At(value = "INVOKE", target = "Lnet/minecraft/server/world/ServerWorld;close()V"), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
 	private void onUnloadWorldAtShutdown(CallbackInfo ci, Iterator<ServerWorld> worlds, ServerWorld world) {
 		ServerWorldEvents.UNLOAD.invoker().onWorldUnload((MinecraftServer) (Object) this, world);
 	}

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/mixin/event/lifecycle/WorldMixin.java

@@ -16,8 +16,12 @@
 
 package net.fabricmc.fabric.mixin.event.lifecycle;
 
+import java.util.HashSet;
+import java.util.Set;
+
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
@@ -25,50 +29,21 @@ import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 import net.minecraft.server.world.ServerWorld;
 import net.minecraft.util.profiler.Profiler;
 import net.minecraft.world.World;
+import net.minecraft.world.chunk.WorldChunk;
 
 import net.fabricmc.fabric.api.event.lifecycle.v1.ServerTickEvents;
+import net.fabricmc.fabric.impl.event.lifecycle.LoadedChunksCache;
 
 @Mixin(World.class)
-public abstract class WorldMixin {
+public abstract class WorldMixin implements LoadedChunksCache {
 	@Shadow
 	public abstract boolean isClient();
 
 	@Shadow
 	public abstract Profiler getProfiler();
 
-	/*@Inject(method = "addBlockEntity", at = @At("TAIL"))
+	@Unique
-	protected void onLoadBlockEntity(BlockEntity blockEntity, CallbackInfoReturnable<Boolean> cir) {
+	private final Set<WorldChunk> loadedChunks = new HashSet<>();
-		if (!this.isClient()) { // Only fire this event if we are a server world
-			ServerBlockEntityEvents.BLOCK_ENTITY_LOAD.invoker().onLoad(blockEntity, (ServerWorld) (Object) this);
-		}
-	}
-
-	// Mojang what hell, why do you need three ways to unload block entities
-	@Inject(method = "removeBlockEntity", at = @At(value = "INVOKE", target = "Ljava/util/List;remove(Ljava/lang/Object;)Z", ordinal = 1), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
-	protected void onUnloadBlockEntity(BlockPos pos, CallbackInfo ci, BlockEntity blockEntity) {
-		if (!this.isClient()) { // Only fire this event if we are a server world
-			ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, (ServerWorld) (Object) this);
-		}
-	}
-
-	@Inject(method = "tickBlockEntities", at = @At(value = "INVOKE", target = "Ljava/util/List;remove(Ljava/lang/Object;)Z"), slice = @Slice(from = @At(value = "INVOKE", target = "Lnet/minecraft/util/profiler/Profiler;pop()V"), to = @At(value = "INVOKE", target = "Lnet/minecraft/world/chunk/WorldChunk;removeBlockEntity(Lnet/minecraft/util/math/BlockPos;)V")), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
-	protected void onRemoveBlockEntity(CallbackInfo ci, Profiler profiler, Iterator iterator, BlockEntity blockEntity) {
-		if (!this.isClient()) {
-			ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, (ServerWorld) (Object) this);
-		}
-	}
-
-	@Redirect(method = "tickBlockEntities", at = @At(value = "INVOKE", target = "Ljava/util/List;removeAll(Ljava/util/Collection;)Z", ordinal = 1))
-	protected boolean onPurgeRemovedBlockEntities(List<BlockEntity> blockEntityList, Collection<BlockEntity> removals) {
-		if (!this.isClient()) {
-			for (BlockEntity removal : removals) {
-				ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(removal, (ServerWorld) (Object) this);
-			}
-		}
-
-		// Mimic vanilla logic
-		return blockEntityList.removeAll(removals);
-	}*/
 
 	@Inject(at = @At("RETURN"), method = "tickBlockEntities")
 	protected void tickWorldAfterBlockEntities(CallbackInfo ci) {
@@ -76,4 +51,19 @@ public abstract class WorldMixin {
 			ServerTickEvents.END_WORLD_TICK.invoker().onEndTick((ServerWorld) (Object) this);
 		}
 	}
+
+	@Override
+	public Set<WorldChunk> fabric_getLoadedChunks() {
+		return this.loadedChunks;
+	}
+
+	@Override
+	public void fabric_markLoaded(WorldChunk chunk) {
+		this.loadedChunks.add(chunk);
+	}
+
+	@Override
+	public void fabric_markUnloaded(WorldChunk chunk) {
+		this.loadedChunks.remove(chunk);
+	}
 }

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/mixin/event/lifecycle/client/ClientPlayNetworkHandlerMixin.java

@@ -22,19 +22,23 @@ import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
+import net.minecraft.block.entity.BlockEntity;
 import net.minecraft.client.network.ClientPlayNetworkHandler;
 import net.minecraft.client.world.ClientWorld;
 import net.minecraft.entity.Entity;
 import net.minecraft.network.packet.s2c.play.GameJoinS2CPacket;
 import net.minecraft.network.packet.s2c.play.PlayerRespawnS2CPacket;
+import net.minecraft.world.chunk.WorldChunk;
 
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
+import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientBlockEntityEvents;
 import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientEntityEvents;
+import net.fabricmc.fabric.impl.event.lifecycle.LoadedChunksCache;
 
 @Environment(EnvType.CLIENT)
 @Mixin(ClientPlayNetworkHandler.class)
-public abstract class ClientPlayNetworkHandlerMixin {
+abstract class ClientPlayNetworkHandlerMixin {
 	@Shadow
 	private ClientWorld world;
 
@@ -42,14 +46,15 @@ public abstract class ClientPlayNetworkHandlerMixin {
 	private void onPlayerRespawn(PlayerRespawnS2CPacket packet, CallbackInfo ci) {
 		// If a world already exists, we need to unload all (block)entities in the world.
 		if (this.world != null) {
-			for (Entity entity : world.getEntities()) {
+			for (Entity entity : this.world.getEntities()) {
 				ClientEntityEvents.ENTITY_UNLOAD.invoker().onUnload(entity, this.world);
 			}
 
-			/*for (BlockEntity blockEntity : world.blockEntities) {
+			for (WorldChunk chunk : ((LoadedChunksCache) this.world).fabric_getLoadedChunks()) {
+				for (BlockEntity blockEntity : chunk.getBlockEntities().values()) {
 					ClientBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, this.world);
-				// No need to clear the `tickingBlockEntities` list since it will be null in just an instant
+				}
-			}*/
+			}
 		}
 	}
 
@@ -67,10 +72,11 @@ public abstract class ClientPlayNetworkHandlerMixin {
 				ClientEntityEvents.ENTITY_UNLOAD.invoker().onUnload(entity, this.world);
 			}
 
-			/*for (BlockEntity blockEntity : world.blockEntities) {
+			for (WorldChunk chunk : ((LoadedChunksCache) this.world).fabric_getLoadedChunks()) {
+				for (BlockEntity blockEntity : chunk.getBlockEntities().values()) {
 					ClientBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, this.world);
-				// No need to clear the `tickingBlockEntities` list since it will be null in just an instant
+				}
-			}*/
+			}
 		}
 	}
 
@@ -79,14 +85,15 @@ public abstract class ClientPlayNetworkHandlerMixin {
 	private void onClearWorld(CallbackInfo ci) {
 		// If a world already exists, we need to unload all (block)entities in the world.
 		if (this.world != null) {
-			for (Entity entity : world.getEntities()) {
+			for (Entity entity : this.world.getEntities()) {
 				ClientEntityEvents.ENTITY_UNLOAD.invoker().onUnload(entity, this.world);
 			}
 
-			/*for (BlockEntity blockEntity : world.blockEntities) {
+			for (WorldChunk chunk : ((LoadedChunksCache) this.world).fabric_getLoadedChunks()) {
+				for (BlockEntity blockEntity : chunk.getBlockEntities().values()) {
 					ClientBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, this.world);
-				// No need to clear the `tickingBlockEntities` list since it will be null in just an instant
+				}
-			}*/
+			}
 		}
 	}
 }

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/mixin/event/lifecycle/client/ClientWorldMixin.java

@@ -31,32 +31,6 @@ import net.fabricmc.fabric.mixin.event.lifecycle.WorldMixin;
 @Environment(EnvType.CLIENT)
 @Mixin(ClientWorld.class)
 public abstract class ClientWorldMixin extends WorldMixin {
-	// We override our injection on the clientworld so only the client's block entity invocations will run
-	/*@Override
-	protected void onLoadBlockEntity(BlockEntity blockEntity, CallbackInfoReturnable<Boolean> cir) {
-		ClientBlockEntityEvents.BLOCK_ENTITY_LOAD.invoker().onLoad(blockEntity, (ClientWorld) (Object) this);
-	}
-
-	// We override our injection on the clientworld so only the client's block entity invocations will run
-	@Override
-	protected void onUnloadBlockEntity(BlockPos pos, CallbackInfo ci, BlockEntity blockEntity) {
-		ClientBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, (ClientWorld) (Object) this);
-	}
-
-	@Override
-	protected void onRemoveBlockEntity(CallbackInfo ci, Profiler profiler, Iterator iterator, BlockEntity blockEntity) {
-		ClientBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(blockEntity, (ClientWorld) (Object) this);
-	}
-
-	@Override
-	protected boolean onPurgeRemovedBlockEntities(List<BlockEntity> blockEntityList, Collection<BlockEntity> removals) {
-		for (BlockEntity removal : removals) {
-			ClientBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(removal, (ClientWorld) (Object) this);
-		}
-
-		return super.onPurgeRemovedBlockEntities(blockEntityList, removals); // Call super
-	}*/
-
 	// We override our injection on the clientworld so only the client world's tick invocations will run
 	@Override
 	protected void tickWorldAfterBlockEntities(CallbackInfo ci) {

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/mixin/event/lifecycle/client/WorldChunkMixin.java

@@ -0,0 +1,112 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.mixin.event.lifecycle.client;
+
+import java.util.Map;
+
+import org.jetbrains.annotations.Nullable;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+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 org.spongepowered.asm.mixin.injection.callback.LocalCapture;
+
+import net.minecraft.block.entity.BlockEntity;
+import net.minecraft.client.world.ClientWorld;
+import net.minecraft.server.world.ServerWorld;
+import net.minecraft.util.math.BlockPos;
+import net.minecraft.world.World;
+import net.minecraft.world.chunk.WorldChunk;
+
+import net.fabricmc.api.EnvType;
+import net.fabricmc.api.Environment;
+import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientBlockEntityEvents;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerBlockEntityEvents;
+
+@Environment(EnvType.CLIENT)
+@Mixin(WorldChunk.class)
+abstract class WorldChunkMixin {
+	@Shadow
+	public abstract World getWorld();
+
+	/*
+	 * @Inject(method = "setBlockEntity", at = @At(value = "CONSTANT", args = "nullValue=true"), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+	 *
+	 * i509VCB: Yes this is very brittle.
+	 * Sadly mixin does not want to cooperate with the Inject annotation commented out above.
+	 * Our goal is to place the inject JUST after the possibly removed block entity is stored onto the stack so we can use local capture:
+	 *
+	 *  INVOKEVIRTUAL net/minecraft/util/math/BlockPos.toImmutable ()Lnet/minecraft/util/math/BlockPos;
+	 *  ALOAD 1
+	 *  INVOKEINTERFACE java/util/Map.put (Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object; (itf)
+	 *  CHECKCAST net/minecraft/block/entity/BlockEntity
+	 *  ASTORE 3
+	 *  <======== HERE
+	 * L6
+	 */
+	@Inject(method = "setBlockEntity", at = @At(value = "INVOKE", target = "Ljava/util/Map;put(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;", shift = At.Shift.BY, by = 3), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+	private void onLoadBlockEntity(BlockEntity blockEntity, CallbackInfo ci, BlockPos blockPos, @Nullable BlockEntity removedBlockEntity) {
+		// Only fire the load event if the block entity has actually changed
+		if (blockEntity != null && blockEntity != removedBlockEntity) {
+			if (this.getWorld() instanceof ServerWorld) {
+				ServerBlockEntityEvents.BLOCK_ENTITY_LOAD.invoker().onLoad(blockEntity, (ServerWorld) this.getWorld());
+			} else if (this.getWorld() instanceof ClientWorld) {
+				ClientBlockEntityEvents.BLOCK_ENTITY_LOAD.invoker().onLoad(blockEntity, (ClientWorld) this.getWorld());
+			}
+		}
+	}
+
+	@Inject(method = "setBlockEntity", at = @At(value = "INVOKE", target = "Lnet/minecraft/block/entity/BlockEntity;markRemoved()V", shift = At.Shift.AFTER), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+	private void onRemoveBlockEntity(BlockEntity blockEntity, CallbackInfo info, BlockPos blockPos, @Nullable BlockEntity removedBlockEntity) {
+		if (removedBlockEntity != null) {
+			if (this.getWorld() instanceof ServerWorld) {
+				ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(removedBlockEntity, (ServerWorld) this.getWorld());
+			} else if (this.getWorld() instanceof ClientWorld) {
+				ClientBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(removedBlockEntity, (ClientWorld) this.getWorld());
+			}
+		}
+	}
+
+	@Redirect(method = "getBlockEntity(Lnet/minecraft/util/math/BlockPos;Lnet/minecraft/world/chunk/WorldChunk$CreationType;)Lnet/minecraft/block/entity/BlockEntity;", at = @At(value = "INVOKE", target = "Ljava/util/Map;remove(Ljava/lang/Object;)Ljava/lang/Object;"))
+	private <K, V> Object onRemoveBlockEntity(Map<K, V> map, K key) {
+		@Nullable
+		final V removed = map.remove(key);
+
+		if (removed != null) {
+			if (this.getWorld() instanceof ServerWorld) {
+				ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload((BlockEntity) removed, (ServerWorld) this.getWorld());
+			} else if (this.getWorld() instanceof ClientWorld) {
+				ClientBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload((BlockEntity) removed, (ClientWorld) this.getWorld());
+			}
+		}
+
+		return removed;
+	}
+
+	@Inject(method = "removeBlockEntity", at = @At(value = "INVOKE", target = "Lnet/minecraft/block/entity/BlockEntity;markRemoved()V"), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+	private void onRemoveBlockEntity(BlockPos pos, CallbackInfo ci, @Nullable BlockEntity removed) {
+		if (removed != null) {
+			if (this.getWorld() instanceof ServerWorld) {
+				ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(removed, (ServerWorld) this.getWorld());
+			} else if (this.getWorld() instanceof ClientWorld) {
+				ClientBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(removed, (ClientWorld) this.getWorld());
+			}
+		}
+	}
+}

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/mixin/event/lifecycle/server/WorldChunkMixin.java

@@ -0,0 +1,102 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.mixin.event.lifecycle.server;
+
+import java.util.Map;
+
+import org.jetbrains.annotations.Nullable;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+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 org.spongepowered.asm.mixin.injection.callback.LocalCapture;
+
+import net.minecraft.block.entity.BlockEntity;
+import net.minecraft.server.world.ServerWorld;
+import net.minecraft.util.math.BlockPos;
+import net.minecraft.world.World;
+import net.minecraft.world.chunk.WorldChunk;
+
+import net.fabricmc.api.EnvType;
+import net.fabricmc.api.Environment;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerBlockEntityEvents;
+
+/**
+ * This is a server only mixin for good reason:
+ * Since all block entity tracking is now on the world chunk, we inject into WorldChunk.
+ * In order to prevent client logic from being loaded due to the mixin, we have a mixin for the client and this one for the server.
+ */
+@Environment(EnvType.SERVER)
+@Mixin(WorldChunk.class)
+abstract class WorldChunkMixin {
+	@Shadow
+	public abstract World getWorld();
+
+	/*
+	 * @Inject(method = "setBlockEntity", at = @At(value = "CONSTANT", args = "nullValue=true"), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+	 *
+	 * i509VCB: Yes this is very brittle.
+	 * Sadly mixin does not want to cooperate with the Inject annotation commented out above.
+	 * Our goal is to place the inject JUST after the possibly removed block entity is stored onto the stack so we can use local capture:
+	 *
+	 *  INVOKEVIRTUAL net/minecraft/util/math/BlockPos.toImmutable ()Lnet/minecraft/util/math/BlockPos;
+	 *  ALOAD 1
+	 *  INVOKEINTERFACE java/util/Map.put (Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object; (itf)
+	 *  CHECKCAST net/minecraft/block/entity/BlockEntity
+	 *  ASTORE 3
+	 *  <======== HERE
+	 * L6
+	 */
+	@Inject(method = "setBlockEntity", at = @At(value = "INVOKE", target = "Ljava/util/Map;put(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;", shift = At.Shift.BY, by = 3), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+	private void onLoadBlockEntity(BlockEntity blockEntity, CallbackInfo ci, BlockPos blockPos, @Nullable BlockEntity removedBlockEntity) {
+		// Only fire the load event if the block entity has actually changed
+		if (blockEntity != null && blockEntity != removedBlockEntity) {
+			if (this.getWorld() instanceof ServerWorld) {
+				ServerBlockEntityEvents.BLOCK_ENTITY_LOAD.invoker().onLoad(blockEntity, (ServerWorld) this.getWorld());
+			}
+		}
+	}
+
+	@Inject(method = "setBlockEntity", at = @At(value = "INVOKE", target = "Lnet/minecraft/block/entity/BlockEntity;markRemoved()V", shift = At.Shift.AFTER), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+	private void onRemoveBlockEntity(BlockEntity blockEntity, CallbackInfo info, BlockPos blockPos, @Nullable BlockEntity removedBlockEntity) {
+		if (removedBlockEntity != null) {
+			if (this.getWorld() instanceof ServerWorld) {
+				ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(removedBlockEntity, (ServerWorld) this.getWorld());
+			}
+		}
+	}
+
+	@Redirect(method = "getBlockEntity(Lnet/minecraft/util/math/BlockPos;Lnet/minecraft/world/chunk/WorldChunk$CreationType;)Lnet/minecraft/block/entity/BlockEntity;", at = @At(value = "INVOKE", target = "Ljava/util/Map;remove(Ljava/lang/Object;)Ljava/lang/Object;"))
+	private <K, V> Object onRemoveBlockEntity(Map<K, V> map, K key) {
+		@Nullable final V removed = map.remove(key);
+
+		if (removed != null && this.getWorld() instanceof ServerWorld) {
+			ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload((BlockEntity) removed, (ServerWorld) this.getWorld());
+		}
+
+		return removed;
+	}
+
+	@Inject(method = "removeBlockEntity", at = @At(value = "INVOKE", target = "Lnet/minecraft/block/entity/BlockEntity;markRemoved()V"), locals = LocalCapture.CAPTURE_FAILEXCEPTION)
+	private void onRemoveBlockEntity(BlockPos pos, CallbackInfo ci, @Nullable BlockEntity removed) {
+		if (removed != null && this.getWorld() instanceof ServerWorld) {
+			ServerBlockEntityEvents.BLOCK_ENTITY_UNLOAD.invoker().onUnload(removed, (ServerWorld) this.getWorld());
+		}
+	}
+}

fabric-lifecycle-events-v1/src/main/resources/fabric-lifecycle-events-v1.mixins.json

@@ -12,9 +12,13 @@
   "client": [
     "client.ClientChunkManagerMixin",
     "client.ClientPlayNetworkHandlerMixin",
-    "client.ClientWorldMixin",
     "client.ClientWorldEntityLoaderMixin",
-    "client.MinecraftClientMixin"
+    "client.ClientWorldMixin",
+    "client.MinecraftClientMixin",
+    "client.WorldChunkMixin"
+  ],
+  "server": [
+    "server.WorldChunkMixin"
   ],
   "injectors": {
     "defaultRequire": 1

fabric-lifecycle-events-v1/src/main/resources/fabric.mod.json

@@ -15,6 +15,14 @@
   "authors": [
     "FabricMC"
   ],
+  "entrypoints": {
+    "main": [
+      "net.fabricmc.fabric.impl.event.lifecycle.LifecycleEventsImpl"
+    ],
+    "client": [
+      "net.fabricmc.fabric.impl.event.lifecycle.ClientLifecycleEventsImpl"
+    ]
+  },
   "mixins": [
     "fabric-lifecycle-events-v1.mixins.json"
   ],

fabric-lifecycle-events-v1/src/testmod/java/net/fabricmc/fabric/test/event/lifecycle/ServerBlockEntityLifecycleTests.java

@@ -19,17 +19,26 @@ package net.fabricmc.fabric.test.event.lifecycle;
 import java.util.ArrayList;
 import java.util.List;
 
+import org.apache.logging.log4j.Logger;
+
 import net.minecraft.block.entity.BlockEntity;
+import net.minecraft.server.world.ServerWorld;
+import net.minecraft.util.registry.Registry;
+import net.minecraft.world.chunk.WorldChunk;
 
 import net.fabricmc.api.ModInitializer;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerBlockEntityEvents;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerLifecycleEvents;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerTickEvents;
+import net.fabricmc.fabric.impl.event.lifecycle.LoadedChunksCache;
 
-public class ServerBlockEntityLifecycleTests implements ModInitializer {
+public final class ServerBlockEntityLifecycleTests implements ModInitializer {
-	private static boolean PRINT_SERVER_BLOCKENTITY_MESSAGES = System.getProperty("fabric-lifecycle-events-testmod.printServerBlockEntityMessages") != null;
+	private static final boolean PRINT_SERVER_BLOCKENTITY_MESSAGES = System.getProperty("fabric-lifecycle-events-testmod.printServerBlockEntityMessages") != null;
-	private List<BlockEntity> serverBlockEntities = new ArrayList<>();
+	private final List<BlockEntity> serverBlockEntities = new ArrayList<>();
 
 	@Override
 	public void onInitialize() {
-		/*final Logger logger = ServerLifecycleTests.LOGGER;
+		final Logger logger = ServerLifecycleTests.LOGGER;
 
 		ServerBlockEntityEvents.BLOCK_ENTITY_LOAD.register((blockEntity, world) -> {
 			this.serverBlockEntities.add(blockEntity);
@@ -56,7 +65,11 @@ public class ServerBlockEntityLifecycleTests implements ModInitializer {
 				}
 
 				for (ServerWorld world : minecraftServer.getWorlds()) {
-					int worldEntities = world.blockEntities.size();
+					int worldEntities = 0;
+
+					for (WorldChunk chunk : ((LoadedChunksCache) world).fabric_getLoadedChunks()) {
+						worldEntities += chunk.getBlockEntities().size();
+					}
 
 					if (PRINT_SERVER_BLOCKENTITY_MESSAGES) {
 						logger.info("[SERVER] Tracked BlockEntities in " + world.getRegistryKey().toString() + " - " + worldEntities);
@@ -82,6 +95,6 @@ public class ServerBlockEntityLifecycleTests implements ModInitializer {
 			if (this.serverBlockEntities.size() != 0) {
 				logger.error("[SERVER] Mismatch in tracked blockentities, expected 0");
 			}
-		});*/
+		});
 	}
 }

fabric-lifecycle-events-v1/src/testmod/java/net/fabricmc/fabric/test/event/lifecycle/ServerEntityLifecycleTests.java

@@ -29,9 +29,9 @@ import net.fabricmc.fabric.api.event.lifecycle.v1.ServerEntityEvents;
 /**
  * Tests related to the lifecycle of entities.
  */
-public class ServerEntityLifecycleTests implements ModInitializer {
+public final class ServerEntityLifecycleTests implements ModInitializer {
-	private static boolean PRINT_SERVER_ENTITY_MESSAGES = System.getProperty("fabric-lifecycle-events-testmod.printServerEntityMessages") != null;
+	private static final boolean PRINT_SERVER_ENTITY_MESSAGES = System.getProperty("fabric-lifecycle-events-testmod.printServerEntityMessages") != null;
-	private List<Entity> serverEntities = new ArrayList<>();
+	private final List<Entity> serverEntities = new ArrayList<>();
 
 	@Override
 	public void onInitialize() {

fabric-lifecycle-events-v1/src/testmod/java/net/fabricmc/fabric/test/event/lifecycle/ServerLifecycleTests.java

@@ -26,7 +26,7 @@ import net.fabricmc.fabric.api.event.lifecycle.v1.ServerWorldEvents;
 /**
  * Tests related to the lifecycle of a server.
  */
-public class ServerLifecycleTests implements ModInitializer {
+public final class ServerLifecycleTests implements ModInitializer {
 	public static final Logger LOGGER = LogManager.getLogger("LifecycleEventsTest");
 
 	@Override

fabric-lifecycle-events-v1/src/testmod/java/net/fabricmc/fabric/test/event/lifecycle/ServerResourceReloadTests.java

@@ -22,7 +22,7 @@ import org.apache.logging.log4j.Logger;
 import net.fabricmc.api.ModInitializer;
 import net.fabricmc.fabric.api.event.lifecycle.v1.ServerLifecycleEvents;
 
-public class ServerResourceReloadTests implements ModInitializer {
+public final class ServerResourceReloadTests implements ModInitializer {
 	public static final Logger LOGGER = LogManager.getLogger("LifecycleEventsTest");
 
 	@Override

fabric-lifecycle-events-v1/src/testmod/java/net/fabricmc/fabric/test/event/lifecycle/ServerTickTests.java

@@ -28,8 +28,8 @@ import net.fabricmc.fabric.api.event.lifecycle.v1.ServerTickEvents;
 /**
  * Test related to ticking events on the server.
  */
-public class ServerTickTests implements ModInitializer {
+public final class ServerTickTests implements ModInitializer {
-	private Map<RegistryKey<World>, Integer> tickTracker = new HashMap<>();
+	private final Map<RegistryKey<World>, Integer> tickTracker = new HashMap<>();
 
 	@Override
 	public void onInitialize() {

fabric-lifecycle-events-v1/src/testmod/java/net/fabricmc/fabric/test/event/lifecycle/client/ClientBlockEntityLifecycleTests.java

@@ -19,18 +19,27 @@ package net.fabricmc.fabric.test.event.lifecycle.client;
 import java.util.ArrayList;
 import java.util.List;
 
+import org.apache.logging.log4j.Logger;
+
 import net.minecraft.block.entity.BlockEntity;
+import net.minecraft.util.registry.Registry;
+import net.minecraft.world.chunk.WorldChunk;
 
 import net.fabricmc.api.ClientModInitializer;
+import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientBlockEntityEvents;
+import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientTickEvents;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerLifecycleEvents;
+import net.fabricmc.fabric.impl.event.lifecycle.LoadedChunksCache;
+import net.fabricmc.fabric.test.event.lifecycle.ServerLifecycleTests;
 
-public class ClientBlockEntityLifecycleTests implements ClientModInitializer {
+public final class ClientBlockEntityLifecycleTests implements ClientModInitializer {
-	private static boolean PRINT_CLIENT_BLOCKENTITY_MESSAGES = System.getProperty("fabric-lifecycle-events-testmod.printClientBlockEntityMessages") != null;
+	private static final boolean PRINT_CLIENT_BLOCKENTITY_MESSAGES = System.getProperty("fabric-lifecycle-events-testmod.printClientBlockEntityMessages") != null;
-	private List<BlockEntity> clientBlockEntities = new ArrayList<>();
+	private final List<BlockEntity> clientBlockEntities = new ArrayList<>();
 	private int clientTicks;
 
 	@Override
 	public void onInitializeClient() {
-		/*final Logger logger = ServerLifecycleTests.LOGGER;
+		final Logger logger = ServerLifecycleTests.LOGGER;
 
 		ClientBlockEntityEvents.BLOCK_ENTITY_LOAD.register((blockEntity, world) -> {
 			this.clientBlockEntities.add(blockEntity);
@@ -50,11 +59,16 @@ public class ClientBlockEntityLifecycleTests implements ClientModInitializer {
 
 		ClientTickEvents.END_CLIENT_TICK.register(client -> {
 			if (this.clientTicks++ % 200 == 0 && client.world != null) {
-				final int blockEntities = client.world.blockEntities.size();
+				int blockEntities = 0;
 
 				if (PRINT_CLIENT_BLOCKENTITY_MESSAGES) {
 					logger.info("[CLIENT] Tracked BlockEntities:" + this.clientBlockEntities.size() + " Ticked at: " + this.clientTicks + "ticks");
-					logger.info("[CLIENT] Actual BlockEntities: " + client.world.blockEntities.size());
+
+					for (WorldChunk chunk : ((LoadedChunksCache) client.world).fabric_getLoadedChunks()) {
+						blockEntities += chunk.getBlockEntities().size();
+					}
+
+					logger.info("[CLIENT] Actual BlockEntities: " + blockEntities);
 				}
 
 				if (blockEntities != this.clientBlockEntities.size()) {
@@ -66,13 +80,13 @@ public class ClientBlockEntityLifecycleTests implements ClientModInitializer {
 		});
 
 		ServerLifecycleEvents.SERVER_STOPPED.register(minecraftServer -> {
-			if (!minecraftServer.isDedicated()) { // fixme: Use ClientNetworking#PLAY_DISCONNECTED instead of the server stop callback for testing.
+			if (!minecraftServer.isDedicated()) { // fixme: Use ClientPlayConnectionEvents#DISCONNECT instead of the server stop callback for testing.
 				logger.info("[CLIENT] Disconnected. Tracking: " + this.clientBlockEntities.size() + " blockentities");
 
 				if (this.clientBlockEntities.size() != 0) {
 					logger.error("[CLIENT] Mismatch in tracked blockentities, expected 0");
 				}
 			}
-		});*/
+		});
 	}
 }

fabric-lifecycle-events-v1/src/testmod/java/net/fabricmc/fabric/test/event/lifecycle/client/ClientEntityLifecycleTests.java

@@ -37,9 +37,9 @@ import net.fabricmc.fabric.test.event.lifecycle.ServerLifecycleTests;
  * Tests related to the lifecycle of entities.
  */
 @Environment(EnvType.CLIENT)
-public class ClientEntityLifecycleTests implements ClientModInitializer {
+public final class ClientEntityLifecycleTests implements ClientModInitializer {
-	private static boolean PRINT_CLIENT_ENTITY_MESSAGES = System.getProperty("fabric-lifecycle-events-testmod.printClientEntityMessages") != null;
+	private static final boolean PRINT_CLIENT_ENTITY_MESSAGES = System.getProperty("fabric-lifecycle-events-testmod.printClientEntityMessages") != null;
-	private List<Entity> clientEntities = new ArrayList<>();
+	private final List<Entity> clientEntities = new ArrayList<>();
 	private int clientTicks;
 
 	@Override

fabric-lifecycle-events-v1/src/testmod/java/net/fabricmc/fabric/test/event/lifecycle/client/ClientLifecycleTests.java

@@ -22,7 +22,7 @@ import net.fabricmc.api.Environment;
 import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientLifecycleEvents;
 
 @Environment(EnvType.CLIENT)
-public class ClientLifecycleTests implements ClientModInitializer {
+public final class ClientLifecycleTests implements ClientModInitializer {
 	@Override
 	public void onInitializeClient() {
 		ClientLifecycleEvents.CLIENT_STARTED.register(client -> {

fabric-lifecycle-events-v1/src/testmod/java/net/fabricmc/fabric/test/event/lifecycle/client/ClientTickTests.java

@@ -29,8 +29,8 @@ import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientTickEvents;
 import net.fabricmc.fabric.test.event.lifecycle.ServerLifecycleTests;
 
 @Environment(EnvType.CLIENT)
-public class ClientTickTests implements ClientModInitializer {
+public final class ClientTickTests implements ClientModInitializer {
-	private Map<RegistryKey<World>, Integer> tickTracker = new HashMap<>();
+	private final Map<RegistryKey<World>, Integer> tickTracker = new HashMap<>();
 	private int ticks;
 
 	@Override