From 5c95d19da93bc0fb23c50b8544259eaa11281955 Mon Sep 17 00:00:00 2001 From: BlueAgent Date: Fri, 29 Dec 2023 23:07:05 +0800 Subject: [PATCH] Fix Memory Leak of Air Capability Listeners (#1275) * Invalidate and revive air capabilities * Reuse air capability invalidation listeners * Fix returning null instead of empty optional * Check that the invalidated listener matches current Also implement the cached listener for RegulatorModule as well. * Add missing else branch to reset the cached air capabilities * Split into getting cached and current neighbour air capabilities * move tube module neighbour caching to separate interface reduce code duplication * Implement a cache for generic capabilities and neighbouring capabilities * Stop creating a lambda on every cached capability get operation --------- Co-authored-by: Des Herriott --- .../AbstractAirHandlingBlockEntity.java | 15 +++- .../block/entity/TubeJunctionBlockEntity.java | 15 +++- .../block/entity/VacuumPumpBlockEntity.java | 15 +++- .../common/capabilities/CapabilityCache.java | 56 ++++++++++++++ .../capabilities/MachineAirHandler.java | 32 +++----- .../NeighbouringCapabilityCache.java | 74 +++++++++++++++++++ .../common/tubemodules/RegulatorModule.java | 24 +++--- .../common/tubemodules/VacuumModule.java | 28 +++---- .../common/util/CapabilityUtils.java | 38 ++++++++++ 9 files changed, 246 insertions(+), 51 deletions(-) create mode 100644 src/main/java/me/desht/pneumaticcraft/common/capabilities/CapabilityCache.java create mode 100644 src/main/java/me/desht/pneumaticcraft/common/capabilities/NeighbouringCapabilityCache.java create mode 100644 src/main/java/me/desht/pneumaticcraft/common/util/CapabilityUtils.java diff --git a/src/main/java/me/desht/pneumaticcraft/common/block/entity/AbstractAirHandlingBlockEntity.java b/src/main/java/me/desht/pneumaticcraft/common/block/entity/AbstractAirHandlingBlockEntity.java index c6702e11f..f3fdaf1e3 100644 --- a/src/main/java/me/desht/pneumaticcraft/common/block/entity/AbstractAirHandlingBlockEntity.java +++ b/src/main/java/me/desht/pneumaticcraft/common/block/entity/AbstractAirHandlingBlockEntity.java @@ -48,7 +48,7 @@ public abstract class AbstractAirHandlingBlockEntity extends AbstractTickingBlockEntity { @GuiSynced protected final IAirHandlerMachine airHandler; - private final LazyOptional airHandlerCap; + private LazyOptional airHandlerCap; private final Map> airHandlerMap = new IdentityHashMap<>(); public AbstractAirHandlingBlockEntity(BlockEntityType type, BlockPos pos, BlockState state, PressureTier pressureTier, int volume, int upgradeSlots) { @@ -58,6 +58,19 @@ public AbstractAirHandlingBlockEntity(BlockEntityType type, BlockPos pos, Blo this.airHandlerCap = LazyOptional.of(() -> airHandler); } + @Override + public void invalidateCaps() { + this.airHandlerCap.invalidate(); + this.airHandlerCap = LazyOptional.empty(); + super.invalidateCaps(); + } + + @Override + public void reviveCaps() { + super.reviveCaps(); + this.airHandlerCap = LazyOptional.of(() -> airHandler); + } + @Override public void handleUpdateTag(CompoundTag tag) { super.handleUpdateTag(tag); diff --git a/src/main/java/me/desht/pneumaticcraft/common/block/entity/TubeJunctionBlockEntity.java b/src/main/java/me/desht/pneumaticcraft/common/block/entity/TubeJunctionBlockEntity.java index b107a3fdb..6994a3c6d 100644 --- a/src/main/java/me/desht/pneumaticcraft/common/block/entity/TubeJunctionBlockEntity.java +++ b/src/main/java/me/desht/pneumaticcraft/common/block/entity/TubeJunctionBlockEntity.java @@ -19,7 +19,7 @@ public class TubeJunctionBlockEntity extends AbstractAirHandlingBlockEntity { private final IAirHandlerMachine tube2Handler; - private final LazyOptional tube2Cap; + private LazyOptional tube2Cap; public TubeJunctionBlockEntity(BlockPos pPos, BlockState pState) { super(ModBlockEntities.TUBE_JUNCTION.get(), pPos, pState, PressureTier.TIER_TWO, 4000, 0); @@ -28,6 +28,19 @@ public TubeJunctionBlockEntity(BlockPos pPos, BlockState pState) { this.tube2Cap = LazyOptional.of(() -> tube2Handler); } + @Override + public void invalidateCaps() { + this.tube2Cap.invalidate(); + this.tube2Cap = LazyOptional.empty(); + super.invalidateCaps(); + } + + @Override + public void reviveCaps() { + super.reviveCaps(); + this.tube2Cap = LazyOptional.of(() -> tube2Handler); + } + @Override public IItemHandler getPrimaryInventory() { return null; diff --git a/src/main/java/me/desht/pneumaticcraft/common/block/entity/VacuumPumpBlockEntity.java b/src/main/java/me/desht/pneumaticcraft/common/block/entity/VacuumPumpBlockEntity.java index b8bdb3c47..433129933 100644 --- a/src/main/java/me/desht/pneumaticcraft/common/block/entity/VacuumPumpBlockEntity.java +++ b/src/main/java/me/desht/pneumaticcraft/common/block/entity/VacuumPumpBlockEntity.java @@ -54,7 +54,7 @@ public class VacuumPumpBlockEntity extends AbstractAirHandlingBlockEntity implem IRedstoneControl, IManoMeasurable, MenuProvider { @GuiSynced private final IAirHandlerMachine vacuumHandler; - private final LazyOptional vacuumCap; + private LazyOptional vacuumCap; public int rotation; public int oldRotation; private int turnTimer = -1; @@ -71,6 +71,19 @@ public VacuumPumpBlockEntity(BlockPos pos, BlockState state) { this.vacuumCap = LazyOptional.of(() -> vacuumHandler); } + @Override + public void invalidateCaps() { + this.vacuumCap.invalidate(); + this.vacuumCap = LazyOptional.empty(); + super.invalidateCaps(); + } + + @Override + public void reviveCaps() { + super.reviveCaps(); + this.vacuumCap = LazyOptional.of(() -> vacuumHandler); + } + @Nonnull @Override public LazyOptional getCapability(@Nonnull Capability cap, @Nullable Direction side) { diff --git a/src/main/java/me/desht/pneumaticcraft/common/capabilities/CapabilityCache.java b/src/main/java/me/desht/pneumaticcraft/common/capabilities/CapabilityCache.java new file mode 100644 index 000000000..3509afd92 --- /dev/null +++ b/src/main/java/me/desht/pneumaticcraft/common/capabilities/CapabilityCache.java @@ -0,0 +1,56 @@ +package me.desht.pneumaticcraft.common.capabilities; + +import net.minecraftforge.common.util.LazyOptional; +import net.minecraftforge.common.util.NonNullConsumer; +import org.jetbrains.annotations.NotNull; + +/** + * Generic capability cache. + *
+ * Typical usage is to {@link CapabilityCache#get()} and use it if it is present, otherwise call {@link CapabilityCache#set(LazyOptional)} to update the capability cache. + */ +public final class CapabilityCache { + private LazyOptional cachedCapability; + private final NonNullConsumer> capabilityInvalidationListener; + + public CapabilityCache() { + this.cachedCapability = LazyOptional.empty(); + this.capabilityInvalidationListener = l -> { + if (this.cachedCapability != l) { + return; + } + + this.cachedCapability = LazyOptional.empty(); + }; + } + + public @NotNull LazyOptional get() { + return this.cachedCapability; + } + + /** + * Sets the cached capability. + * Also handles registering an invalidation listener to clear the cached capability. + * + * @param cap Capability to set. + * @return The capability that was set. + */ + public @NotNull LazyOptional set(LazyOptional cap) { + if (this.cachedCapability == cap) { + return cap; + } + + if (!cap.isPresent()) { + this.cachedCapability = LazyOptional.empty(); + return this.cachedCapability; + } + + this.cachedCapability = cap; + cap.addListener(this.capabilityInvalidationListener); + return this.cachedCapability; + } + + public void clear() { + this.cachedCapability = LazyOptional.empty(); + } +} diff --git a/src/main/java/me/desht/pneumaticcraft/common/capabilities/MachineAirHandler.java b/src/main/java/me/desht/pneumaticcraft/common/capabilities/MachineAirHandler.java index 9f649d808..edf3595b0 100644 --- a/src/main/java/me/desht/pneumaticcraft/common/capabilities/MachineAirHandler.java +++ b/src/main/java/me/desht/pneumaticcraft/common/capabilities/MachineAirHandler.java @@ -58,7 +58,7 @@ public class MachineAirHandler extends BasicAirHandler implements IAirHandlerMac private Direction leakDir = null; private Direction prevLeakDir = null; private int prevAir; - private final Map> neighbourAirHandlers = new EnumMap<>(Direction.class); + private final NeighbouringCapabilityCache neighbouringAirHandlerCache; // note: leaks due to security upgrade are tracked separately from leaks due to disconnection private boolean safetyLeaking; // is the handler venting right now? @@ -69,9 +69,7 @@ public MachineAirHandler(PressureTier tier, int volume) { super(volume); this.tier = tier; - for (Direction dir : DirectionUtil.VALUES) { - this.neighbourAirHandlers.put(dir, LazyOptional.empty()); - } + this.neighbouringAirHandlerCache = new NeighbouringCapabilityCache<>(); } @Override @@ -122,10 +120,7 @@ public void setConnectedFaces(List sides) { connectedFaces.clear(); sides.forEach(side -> connectedFaces.set(side.get3DDataValue())); - // invalidate cached neighbour data - for (Direction dir : DirectionUtil.VALUES) { - this.neighbourAirHandlers.put(dir, LazyOptional.empty()); - } + this.neighbouringAirHandlerCache.clear(); } @Override @@ -240,22 +235,15 @@ public Direction getSideLeaking() { return this.leakDir; } - private LazyOptional getNeighbourAirHandler(BlockEntity ownerTE, Direction dir) { + private LazyOptional getCachedNeighbourAirHandler(BlockEntity ownerTE, Direction dir) { if (!connectedFaces.get(dir.get3DDataValue())) return LazyOptional.empty(); - if (!neighbourAirHandlers.get(dir).isPresent()) { - BlockEntity te1 = Objects.requireNonNull(ownerTE.getLevel()).getBlockEntity(ownerTE.getBlockPos().relative(dir)); - if (te1 != null) { - LazyOptional cap = te1.getCapability(PNCCapabilities.AIR_HANDLER_MACHINE_CAPABILITY, dir.getOpposite()); - if (cap.isPresent()) { - neighbourAirHandlers.put(dir, cap); - neighbourAirHandlers.get(dir).addListener(l -> neighbourAirHandlers.put(dir, LazyOptional.empty())); - } - } else { - neighbourAirHandlers.put(dir, LazyOptional.empty()); - } + LazyOptional cap = this.neighbouringAirHandlerCache.get(dir); + if (cap.isPresent()) { + return cap; } - return neighbourAirHandlers.get(dir); + + return this.neighbouringAirHandlerCache.set(PNCCapabilities.AIR_HANDLER_MACHINE_CAPABILITY, ownerTE, dir); } private void disperseAir(BlockEntity ownerTE) { @@ -292,7 +280,7 @@ private List getConnectedAirHandlers(BlockEntity ownerTE, boolean on List neighbours = new ArrayList<>(); for (Direction dir : DirectionUtil.VALUES) { if (connectedFaces.get(dir.get3DDataValue())) { - getNeighbourAirHandler(ownerTE, dir).ifPresent(h -> { + getCachedNeighbourAirHandler(ownerTE, dir).ifPresent(h -> { if ((!onlyLowerPressure || h.getPressure() < getPressure())) { neighbours.add(new ConnectedAirHandler(dir, h)); } diff --git a/src/main/java/me/desht/pneumaticcraft/common/capabilities/NeighbouringCapabilityCache.java b/src/main/java/me/desht/pneumaticcraft/common/capabilities/NeighbouringCapabilityCache.java new file mode 100644 index 000000000..23b718832 --- /dev/null +++ b/src/main/java/me/desht/pneumaticcraft/common/capabilities/NeighbouringCapabilityCache.java @@ -0,0 +1,74 @@ +package me.desht.pneumaticcraft.common.capabilities; + +import me.desht.pneumaticcraft.common.util.CapabilityUtils; +import net.minecraft.core.Direction; +import net.minecraft.world.level.block.entity.BlockEntity; +import net.minecraftforge.common.capabilities.Capability; +import net.minecraftforge.common.util.LazyOptional; +import org.jetbrains.annotations.NotNull; + +import java.util.EnumMap; +import java.util.Map; + +/** + * Capability cache for neighbours of a {@link BlockEntity}. + *
+ * Typical usage is to {@link NeighbouringCapabilityCache#get(Direction)} and use it if it is present, + * otherwise call {@link NeighbouringCapabilityCache#set(Capability, BlockEntity, Direction)} to update the capability cache. + */ +public class NeighbouringCapabilityCache { + private final Map> neighbouringCapabilityCaches; + + public NeighbouringCapabilityCache() { + this.neighbouringCapabilityCaches = new EnumMap<>(Direction.class); + } + + /** + * Get cached capability for a neighbouring {@link BlockEntity}. + * + * @param direction Direction from the original block entity to a neighbouring block entity. + * @return The cached capability. + */ + public @NotNull LazyOptional get(@NotNull Direction direction) { + CapabilityCache cache = this.neighbouringCapabilityCaches.get(direction); + if (cache == null) { + return LazyOptional.empty(); + } + + return cache.get(); + } + + /** + * Sets the cached capability for a neighbouring {@link BlockEntity}. + * Also handles registering an invalidation listener to clear the cached capability. + * + * @param capability Capability to get and cache. + * @param blockEntity Originating {@link BlockEntity}. + * @param direction Direction from {@param blockEntity} to a neighbouring block entity to look for a capability. + * @return The cached capability. + */ + public @NotNull LazyOptional set(@NotNull Capability capability, @NotNull BlockEntity blockEntity, @NotNull Direction direction) { + CapabilityCache cache = this.neighbouringCapabilityCaches.get(direction); + if (cache == null) { + cache = new CapabilityCache<>(); + this.neighbouringCapabilityCaches.put(direction, cache); + } + + return cache.set(CapabilityUtils.getNeighbourCap(capability, blockEntity, direction)); + } + + public void clear() { + for (Direction dir : Direction.values()) { + clear(dir); + } + } + + public void clear(@NotNull Direction dir) { + CapabilityCache cache = this.neighbouringCapabilityCaches.get(dir); + if (cache == null) { + return; + } + + cache.clear(); + } +} diff --git a/src/main/java/me/desht/pneumaticcraft/common/tubemodules/RegulatorModule.java b/src/main/java/me/desht/pneumaticcraft/common/tubemodules/RegulatorModule.java index 574097341..8f2c0e927 100644 --- a/src/main/java/me/desht/pneumaticcraft/common/tubemodules/RegulatorModule.java +++ b/src/main/java/me/desht/pneumaticcraft/common/tubemodules/RegulatorModule.java @@ -21,18 +21,21 @@ import me.desht.pneumaticcraft.api.tileentity.IAirHandlerMachine; import me.desht.pneumaticcraft.common.block.PressureTubeBlock; import me.desht.pneumaticcraft.common.block.entity.PressureTubeBlockEntity; +import me.desht.pneumaticcraft.common.capabilities.CapabilityCache; import me.desht.pneumaticcraft.common.core.ModItems; +import me.desht.pneumaticcraft.common.util.CapabilityUtils; import me.desht.pneumaticcraft.lib.PneumaticValues; import net.minecraft.core.Direction; import net.minecraft.world.item.Item; -import net.minecraft.world.level.block.entity.BlockEntity; import net.minecraftforge.common.util.LazyOptional; public class RegulatorModule extends AbstractRedstoneReceivingModule implements IInfluenceDispersing { - private LazyOptional neighbourCap = null; + private final CapabilityCache neighbourAirHandlerCache; public RegulatorModule(Direction dir, PressureTubeBlockEntity pressureTube) { super(dir, pressureTube); + + this.neighbourAirHandlerCache = new CapabilityCache<>(); } @Override @@ -60,7 +63,8 @@ public boolean isInline() { @Override public void onNeighborBlockUpdate() { super.onNeighborBlockUpdate(); - neighbourCap = null; + + this.neighbourAirHandlerCache.clear(); } @Override @@ -70,16 +74,12 @@ public float getThreshold() { } private LazyOptional getCachedNeighbourAirHandler() { - if (neighbourCap == null) { - BlockEntity neighborTE = pressureTube.nonNullLevel().getBlockEntity(pressureTube.getBlockPos().relative(dir)); - if (neighborTE != null) { - neighbourCap = neighborTE.getCapability(PNCCapabilities.AIR_HANDLER_MACHINE_CAPABILITY, dir.getOpposite()); - if (neighbourCap.isPresent()) neighbourCap.addListener(l -> neighbourCap = null); - } else { - neighbourCap = LazyOptional.empty(); - } + LazyOptional cap = this.neighbourAirHandlerCache.get(); + if (cap.isPresent()) { + return cap; } - return neighbourCap; + + return this.neighbourAirHandlerCache.set(CapabilityUtils.getNeighbourCap(PNCCapabilities.AIR_HANDLER_MACHINE_CAPABILITY, pressureTube, dir)); } @Override diff --git a/src/main/java/me/desht/pneumaticcraft/common/tubemodules/VacuumModule.java b/src/main/java/me/desht/pneumaticcraft/common/tubemodules/VacuumModule.java index 4f1ee6e3c..a9d610949 100644 --- a/src/main/java/me/desht/pneumaticcraft/common/tubemodules/VacuumModule.java +++ b/src/main/java/me/desht/pneumaticcraft/common/tubemodules/VacuumModule.java @@ -21,23 +21,26 @@ import me.desht.pneumaticcraft.api.tileentity.IAirHandlerMachine; import me.desht.pneumaticcraft.common.block.PressureTubeBlock; import me.desht.pneumaticcraft.common.block.entity.PressureTubeBlockEntity; +import me.desht.pneumaticcraft.common.capabilities.CapabilityCache; import me.desht.pneumaticcraft.common.core.ModItems; +import me.desht.pneumaticcraft.common.util.CapabilityUtils; import me.desht.pneumaticcraft.common.util.PneumaticCraftUtils; import me.desht.pneumaticcraft.lib.PneumaticValues; import net.minecraft.core.Direction; import net.minecraft.nbt.CompoundTag; import net.minecraft.world.item.Item; -import net.minecraft.world.level.block.entity.BlockEntity; import net.minecraftforge.common.util.LazyOptional; public class VacuumModule extends AbstractRedstoneReceivingModule implements IInfluenceDispersing { - private LazyOptional neighbourCap = null; + private final CapabilityCache neighbourAirHandlerCache; public float rotation, oldRotation; private int lastAmount = 0; public VacuumModule(Direction dir, PressureTubeBlockEntity pressureTube) { super(dir, pressureTube); + + this.neighbourAirHandlerCache = new CapabilityCache<>(); } @Override @@ -55,7 +58,7 @@ public void tickCommon() { public void onNeighborBlockUpdate() { super.onNeighborBlockUpdate(); - neighbourCap = null; + this.neighbourAirHandlerCache.clear(); } @Override @@ -64,11 +67,12 @@ public void tickServer() { int prevLast = lastAmount; PressureTubeBlockEntity tube = getTube(); - if (tube.getPressure() >= PneumaticValues.MIN_PRESSURE_VACUUM_PUMP && getCachedNeighbourAirHandler().isPresent() && getReceivingRedstoneLevel() == 0) { + LazyOptional neighbourCap; + if (tube.getPressure() >= PneumaticValues.MIN_PRESSURE_VACUUM_PUMP && (neighbourCap = getCachedNeighbourAirHandler()).isPresent() && getReceivingRedstoneLevel() == 0) { int toAdd = (int) (-PneumaticValues.USAGE_VACUUM_PUMP * (upgraded ? 7.41f : 1)); int toTake = (int) (-PneumaticValues.PRODUCTION_VACUUM_PUMP * (upgraded ? 5.06f : 1)); - lastAmount = getCachedNeighbourAirHandler().map(h -> { + lastAmount = neighbourCap.map(h -> { int air = h.getAir(); float pressure = h.getPressure(); h.addAir(toTake); @@ -106,16 +110,12 @@ public void setLastAmount(int lastAmount) { } private LazyOptional getCachedNeighbourAirHandler() { - if (neighbourCap == null) { - BlockEntity neighborTE = pressureTube.nonNullLevel().getBlockEntity(pressureTube.getBlockPos().relative(dir)); - if (neighborTE != null) { - neighbourCap = neighborTE.getCapability(PNCCapabilities.AIR_HANDLER_MACHINE_CAPABILITY, dir.getOpposite()); - if (neighbourCap.isPresent()) neighbourCap.addListener(l -> neighbourCap = null); - } else { - neighbourCap = LazyOptional.empty(); - } + LazyOptional cap = this.neighbourAirHandlerCache.get(); + if (cap.isPresent()) { + return cap; } - return neighbourCap; + + return this.neighbourAirHandlerCache.set(CapabilityUtils.getNeighbourCap(PNCCapabilities.AIR_HANDLER_MACHINE_CAPABILITY, pressureTube, dir)); } @Override diff --git a/src/main/java/me/desht/pneumaticcraft/common/util/CapabilityUtils.java b/src/main/java/me/desht/pneumaticcraft/common/util/CapabilityUtils.java new file mode 100644 index 000000000..fc139d95e --- /dev/null +++ b/src/main/java/me/desht/pneumaticcraft/common/util/CapabilityUtils.java @@ -0,0 +1,38 @@ +package me.desht.pneumaticcraft.common.util; + +import net.minecraft.core.Direction; +import net.minecraft.world.level.Level; +import net.minecraft.world.level.block.entity.BlockEntity; +import net.minecraftforge.common.capabilities.Capability; +import net.minecraftforge.common.util.LazyOptional; +import org.jetbrains.annotations.NotNull; + +public final class CapabilityUtils { + private CapabilityUtils() { + } + + /** + * Gets capability of a neighbouring {@link BlockEntity}. + * + * @param capability Capability to get. + * @param blockEntity Originating {@link BlockEntity}. + * @param direction Direction from {@param blockEntity} to a neighbouring block entity to look for a capability. + * @return The capability on the neighbouring {@link BlockEntity}. + */ + public static @NotNull LazyOptional getNeighbourCap(@NotNull Capability capability, @NotNull BlockEntity blockEntity, @NotNull Direction direction) { + Level level = blockEntity.getLevel(); + if (level == null) { + return LazyOptional.empty(); + } + + // We could check if the neighbouring block is loaded to avoid loading unloaded chunks? + // This would be a behaviour change from current though. + + BlockEntity neighbourTE = level.getBlockEntity(blockEntity.getBlockPos().relative(direction)); + if (neighbourTE == null) { + return LazyOptional.empty(); + } + + return neighbourTE.getCapability(capability, direction.getOpposite()); + } +}