Retain pregen dependency chunks

- PregenChunkRetention holds every chunk in a FULL request's dependency square until all in-area regions reaching it drained; drained regions keep only the border strips undrained neighbours need
- Removes the unload/reload churn (about 1.5 reloads per generated chunk at r1000: lattice passes, next region's border, future-region and ring protos)
- Capped at one chunk per MB of heap; over the cap the drained strips farthest from the active region are released; heap pressure and close release everything
- Regions still save and unload at drain; the early-unload test now pins that plus border retention

adapters/bukkit/nms/v26_2_R1/src/test/java/art/arcane/iris/platform/bukkit/nms/v26_2_R1/NmsChunkRetentionTest.java

@@ -0,0 +1,84 @@
+package art.arcane.iris.platform.bukkit.nms.v26_2_R1;
+
+import art.arcane.volmlib.nativelib.v26_2_R1.terrain.NativeTerrainAccessImpl;
+import ca.spottedleaf.moonrise.patches.chunk_system.scheduling.ChunkHolderManager;
+import ca.spottedleaf.moonrise.patches.chunk_system.scheduling.ChunkTaskScheduler;
+import net.minecraft.SharedConstants;
+import net.minecraft.server.Bootstrap;
+import net.minecraft.server.level.ChunkLevel;
+import net.minecraft.server.level.ServerLevel;
+import net.minecraft.server.level.TicketType;
+import net.minecraft.world.level.chunk.status.ChunkStatus;
+import org.bukkit.craftbukkit.CraftWorld;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+
+import java.lang.reflect.Field;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.ArgumentMatchers.same;
+import static org.mockito.Mockito.CALLS_REAL_METHODS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class NmsChunkRetentionTest {
+    @BeforeClass
+    public static void bootstrapMinecraftRegistries() {
+        SharedConstants.tryDetectVersion();
+        Bootstrap.bootStrap();
+    }
+
+    @Test
+    public void retentionRequestsNothingBeyondTheOuterRingOfAFullRequest() {
+        NativeTerrainAccessImpl binding = mock(NativeTerrainAccessImpl.class, CALLS_REAL_METHODS);
+        assertEquals(ChunkLevel.RADIUS_AROUND_FULL_CHUNK, binding.fullChunkDependencyRadius());
+        assertEquals(ChunkLevel.MAX_LEVEL, ChunkHolderManager.MAX_TICKET_LEVEL);
+        assertEquals(ChunkLevel.generationStatus(ChunkLevel.FULL_CHUNK_LEVEL + binding.fullChunkDependencyRadius()),
+                ChunkLevel.generationStatus(ChunkHolderManager.MAX_TICKET_LEVEL));
+        assertEquals(ChunkStatus.STRUCTURE_STARTS, ChunkLevel.generationStatus(ChunkHolderManager.MAX_TICKET_LEVEL));
+    }
+
+    @Test
+    public void retainAndReleaseUseOnePermanentNonPersistentTicketAtTheMaximumLevel() throws Exception {
+        NativeTerrainAccessImpl binding = mock(NativeTerrainAccessImpl.class, CALLS_REAL_METHODS);
+        CraftWorld world = mock(CraftWorld.class);
+        ServerLevel level = mock(ServerLevel.class);
+        ChunkTaskScheduler scheduler = mock(ChunkTaskScheduler.class);
+        ChunkHolderManager manager = mock(ChunkHolderManager.class);
+        Field managerField = ChunkTaskScheduler.class.getField("chunkHolderManager");
+        managerField.setAccessible(true);
+        managerField.set(scheduler, manager);
+        when(world.getHandle()).thenReturn(level);
+        when(level.moonrise$getChunkTaskScheduler()).thenReturn(scheduler);
+        when(manager.addTicketAtLevel(any(TicketType.class), eq(3), eq(-4), eq(ChunkHolderManager.MAX_TICKET_LEVEL), isNull()))
+                .thenReturn(true);
+        when(manager.removeTicketAtLevel(any(TicketType.class), eq(3), eq(-4), eq(ChunkHolderManager.MAX_TICKET_LEVEL), isNull()))
+                .thenReturn(true);
+
+        assertTrue(binding.retainChunk(world, 3, -4));
+        ArgumentCaptor<TicketType> type = ArgumentCaptor.forClass(TicketType.class);
+        verify(manager).addTicketAtLevel(type.capture(), eq(3), eq(-4), eq(ChunkHolderManager.MAX_TICKET_LEVEL), isNull());
+        assertEquals(TicketType.NO_TIMEOUT, type.getValue().timeout());
+        assertTrue(type.getValue().doesLoad());
+        assertFalse(type.getValue().doesSimulate());
+        assertFalse(type.getValue().persist());
+
+        assertTrue(binding.releaseChunk(world, 3, -4));
+        verify(manager).removeTicketAtLevel(same(type.getValue()), eq(3), eq(-4), eq(ChunkHolderManager.MAX_TICKET_LEVEL), isNull());
+
+        binding.retainChunk(world, 3, -4);
+        ArgumentCaptor<TicketType> again = ArgumentCaptor.forClass(TicketType.class);
+        verify(manager, times(2)).addTicketAtLevel(again.capture(), eq(3), eq(-4),
+                eq(ChunkHolderManager.MAX_TICKET_LEVEL), isNull());
+        assertSame(type.getValue(), again.getValue());
+    }
+}

adapters/bukkit/nms/v26_3_R1/src/test/java/art/arcane/iris/platform/bukkit/nms/v26_3_R1/NmsChunkRetentionTest.java

@@ -0,0 +1,84 @@
+package art.arcane.iris.platform.bukkit.nms.v26_3_R1;
+
+import art.arcane.volmlib.nativelib.v26_3_R1.terrain.NativeTerrainAccessImpl;
+import ca.spottedleaf.moonrise.patches.chunk_system.scheduling.ChunkHolderManager;
+import ca.spottedleaf.moonrise.patches.chunk_system.scheduling.ChunkTaskScheduler;
+import net.minecraft.SharedConstants;
+import net.minecraft.server.Bootstrap;
+import net.minecraft.server.level.ChunkLevel;
+import net.minecraft.server.level.ServerLevel;
+import net.minecraft.server.level.TicketType;
+import net.minecraft.world.level.chunk.status.ChunkStatus;
+import org.bukkit.craftbukkit.CraftWorld;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+
+import java.lang.reflect.Field;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.ArgumentMatchers.same;
+import static org.mockito.Mockito.CALLS_REAL_METHODS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class NmsChunkRetentionTest {
+    @BeforeClass
+    public static void bootstrapMinecraftRegistries() {
+        SharedConstants.tryDetectVersion();
+        Bootstrap.bootStrap();
+    }
+
+    @Test
+    public void retentionRequestsNothingBeyondTheOuterRingOfAFullRequest() {
+        NativeTerrainAccessImpl binding = mock(NativeTerrainAccessImpl.class, CALLS_REAL_METHODS);
+        assertEquals(ChunkLevel.RADIUS_AROUND_FULL_CHUNK, binding.fullChunkDependencyRadius());
+        assertEquals(ChunkLevel.MAX_LEVEL, ChunkHolderManager.MAX_TICKET_LEVEL);
+        assertEquals(ChunkLevel.generationStatus(ChunkLevel.FULL_CHUNK_LEVEL + binding.fullChunkDependencyRadius()),
+                ChunkLevel.generationStatus(ChunkHolderManager.MAX_TICKET_LEVEL));
+        assertEquals(ChunkStatus.STRUCTURE_STARTS, ChunkLevel.generationStatus(ChunkHolderManager.MAX_TICKET_LEVEL));
+    }
+
+    @Test
+    public void retainAndReleaseUseOnePermanentNonPersistentTicketAtTheMaximumLevel() throws Exception {
+        NativeTerrainAccessImpl binding = mock(NativeTerrainAccessImpl.class, CALLS_REAL_METHODS);
+        CraftWorld world = mock(CraftWorld.class);
+        ServerLevel level = mock(ServerLevel.class);
+        ChunkTaskScheduler scheduler = mock(ChunkTaskScheduler.class);
+        ChunkHolderManager manager = mock(ChunkHolderManager.class);
+        Field managerField = ChunkTaskScheduler.class.getField("chunkHolderManager");
+        managerField.setAccessible(true);
+        managerField.set(scheduler, manager);
+        when(world.getHandle()).thenReturn(level);
+        when(level.moonrise$getChunkTaskScheduler()).thenReturn(scheduler);
+        when(manager.addTicketAtLevel(any(TicketType.class), eq(3), eq(-4), eq(ChunkHolderManager.MAX_TICKET_LEVEL), isNull()))
+                .thenReturn(true);
+        when(manager.removeTicketAtLevel(any(TicketType.class), eq(3), eq(-4), eq(ChunkHolderManager.MAX_TICKET_LEVEL), isNull()))
+                .thenReturn(true);
+
+        assertTrue(binding.retainChunk(world, 3, -4));
+        ArgumentCaptor<TicketType> type = ArgumentCaptor.forClass(TicketType.class);
+        verify(manager).addTicketAtLevel(type.capture(), eq(3), eq(-4), eq(ChunkHolderManager.MAX_TICKET_LEVEL), isNull());
+        assertEquals(TicketType.NO_TIMEOUT, type.getValue().timeout());
+        assertTrue(type.getValue().doesLoad());
+        assertFalse(type.getValue().doesSimulate());
+        assertFalse(type.getValue().persist());
+
+        assertTrue(binding.releaseChunk(world, 3, -4));
+        verify(manager).removeTicketAtLevel(same(type.getValue()), eq(3), eq(-4), eq(ChunkHolderManager.MAX_TICKET_LEVEL), isNull());
+
+        binding.retainChunk(world, 3, -4);
+        ArgumentCaptor<TicketType> again = ArgumentCaptor.forClass(TicketType.class);
+        verify(manager, times(2)).addTicketAtLevel(again.capture(), eq(3), eq(-4),
+                eq(ChunkHolderManager.MAX_TICKET_LEVEL), isNull());
+        assertSame(type.getValue(), again.getValue());
+    }
+}

core/src/main/java/art/arcane/iris/platform/bukkit/nms/INMSBinding.java

@@ -233,6 +233,18 @@ public interface INMSBinding extends NativeTerrainAccess {
         return false;
     }
 
+    default int fullChunkDependencyRadius() {
+        return -1;
+    }
+
+    default boolean retainChunk(World world, int x, int z) {
+        return false;
+    }
+
+    default boolean releaseChunk(World world, int x, int z) {
+        return false;
+    }
+
     default boolean pollChunkTask(World world) {
         return false;
     }

core/src/main/java/art/arcane/iris/world/pregen/AsyncPregenMethod.java

@@ -69,6 +69,7 @@ public class AsyncPregenMethod implements PregeneratorMethod {
     private static final long MANTLE_CLEANUP_DRAIN_SECONDS = 60L;
     private static final long CHUNK_FLUSH_DRAIN_SECONDS = 30L;
     private static final long ADMISSION_WAIT_BOUND_MS = 500L;
+    private static final long RETAINED_CHUNK_BYTES = 1024L * 1024L;
     private final World world;
     private final IrisRuntimeSchedulerMode runtimeSchedulerMode;
     private final IrisPaperLikeBackendMode paperLikeBackendMode;
@@ -116,6 +117,7 @@ public class AsyncPregenMethod implements PregeneratorMethod {
     private volatile Engine metricsEngine;
     private volatile Mantle cachedMantle;
     private volatile HydrologyTileCache.PregenerationScope hydrologyPrefetchScope;
+    private volatile PregenChunkRetention retention;
     private final PregenMantleBackpressure backpressure;
 
     public AsyncPregenMethod(World world, int unusedThreads) {
@@ -241,6 +243,7 @@ public class AsyncPregenMethod implements PregeneratorMethod {
             if (pending != null && pending.decrementAndGet() == 0) {
                 onRegionDrained(rk);
             } else if (MantleHeapPressure.overHighWater()) {
+                releaseRetention();
                 evictRegion(rk);
             }
         } catch (Throwable e) {
@@ -272,6 +275,66 @@ public class AsyncPregenMethod implements PregeneratorMethod {
         boundsMinRegionZ = minRegionZ;
         boundsMaxRegionX = maxRegionX;
         boundsMaxRegionZ = maxRegionZ;
+        retention = createRetention(minRegionX, minRegionZ, maxRegionX, maxRegionZ);
+    }
+
+    private PregenChunkRetention createRetention(int minRegionX, int minRegionZ, int maxRegionX, int maxRegionZ) {
+        if (foliaRuntime) {
+            return null;
+        }
+        try {
+            INMSBinding binding = INMS.get();
+            int radius = binding.fullChunkDependencyRadius();
+            if (radius < 0 || radius > 31) {
+                return null;
+            }
+            PregenChunkRetention.Tickets tickets = new PregenChunkRetention.Tickets() {
+                @Override
+                public void retain(int chunkX, int chunkZ) {
+                    binding.retainChunk(world, chunkX, chunkZ);
+                }
+
+                @Override
+                public void release(int chunkX, int chunkZ) {
+                    binding.releaseChunk(world, chunkX, chunkZ);
+                }
+            };
+            return new PregenChunkRetention(tickets, radius, retentionCapacity(Runtime.getRuntime().maxMemory()),
+                    minRegionX, minRegionZ, maxRegionX, maxRegionZ);
+        } catch (Throwable e) {
+            PregenDiagnostics.probeFailed("pregen chunk retention for world " + world.getName(), e);
+            return null;
+        }
+    }
+
+    static int retentionCapacity(long maximumHeapBytes) {
+        return (int) Math.min(Integer.MAX_VALUE, Math.max(0L, maximumHeapBytes) / RETAINED_CHUNK_BYTES);
+    }
+
+    private void retainAround(int x, int z) {
+        PregenChunkRetention active = retention;
+        if (active == null) {
+            return;
+        }
+        try {
+            active.retainAround(x, z);
+        } catch (Throwable e) {
+            retention = null;
+            IrisLogging.reportError("Async pregen chunk retention failed at " + x + "," + z + " in world " + world.getName()
+                    + "; neighbouring regions reload their borders from disk from now on.", e);
+            try {
+                active.releaseAll();
+            } catch (Throwable releaseFailure) {
+                IrisLogging.reportError(releaseFailure);
+            }
+        }
+    }
+
+    private void releaseRetention() {
+        PregenChunkRetention active = retention;
+        if (active != null) {
+            active.releaseAll();
+        }
     }
 
     @Override
@@ -309,6 +372,10 @@ public class AsyncPregenMethod implements PregeneratorMethod {
 
     private void onRegionDrained(long rk) {
         regionPending.remove(rk);
+        PregenChunkRetention active = retention;
+        if (active != null) {
+            active.drained((int) (rk >> 32), (int) rk);
+        }
         evictRegion(rk);
     }
 
@@ -781,6 +848,7 @@ public class AsyncPregenMethod implements PregeneratorMethod {
             syncGenerationHistory();
 
             mantleCleanup.close(MANTLE_CLEANUP_DRAIN_SECONDS, TimeUnit.SECONDS);
+            releaseRetention();
             flushAllRemainingChunks();
             chunkFlush.close(CHUNK_FLUSH_DRAIN_SECONDS, TimeUnit.SECONDS);
             executor.shutdown();
@@ -830,6 +898,7 @@ public class AsyncPregenMethod implements PregeneratorMethod {
 
     @Override
     public void reclaimMemory() {
+        releaseRetention();
         if (regionChunks.isEmpty() && pendingEvictions.isEmpty()) {
             return;
         }
@@ -879,6 +948,7 @@ public class AsyncPregenMethod implements PregeneratorMethod {
         }
 
         regionPending.computeIfAbsent(rkey(x >> 5, z >> 5), k -> new AtomicInteger(1)).incrementAndGet();
+        retainAround(x, z);
         markSubmitted();
         executor.generate(x, z, listener);
     }

core/src/main/java/art/arcane/iris/world/pregen/PregenChunkRetention.java

@@ -0,0 +1,193 @@
+/*
+ * Iris is a World Generator for Minecraft Bukkit Servers
+ * Copyright (c) 2022 Arcane Arts (Volmit Software)
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package art.arcane.iris.world.pregen;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Keeps the chunks a pregen FULL request loads around itself resident until no later request of the
+ * pregen area can need them. Without it the server unloads a chunk as soon as the in-flight requests
+ * around it move on, and the next request that reaches it (the next lattice pass over the region, or the
+ * neighbouring region's border) deserializes it from disk again.
+ *
+ * <p>A request needs every chunk within the dependency radius, so a retained chunk is needed by exactly
+ * the in-area regions its own radius square overlaps. It is released once all of them have drained: a
+ * drained region keeps only the border strips its undrained neighbours still reach. Retained positions
+ * are capped; over the cap the strips of the drained regions farthest from the active region are released
+ * first (they are needed last), and a released strip is simply reloaded by the request that reaches it, as
+ * the pregen did before retention.
+ */
+final class PregenChunkRetention {
+    private static final int REGION_WORDS = 16;
+
+    private final Tickets tickets;
+    private final int radius;
+    private final int capacity;
+    private final Map<Long, long[]> held = new HashMap<>();
+    private final Set<Long> drained = new HashSet<>();
+    private final int minRegionX;
+    private final int minRegionZ;
+    private final int maxRegionX;
+    private final int maxRegionZ;
+    private int size;
+
+    PregenChunkRetention(Tickets tickets, int radius, int capacity,
+                         int minRegionX, int minRegionZ, int maxRegionX, int maxRegionZ) {
+        if (radius < 0 || radius > 31) {
+            throw new IllegalArgumentException("Retention radius must fit inside one region: " + radius);
+        }
+        this.tickets = tickets;
+        this.radius = radius;
+        this.capacity = Math.max(1, capacity);
+        this.minRegionX = minRegionX;
+        this.minRegionZ = minRegionZ;
+        this.maxRegionX = maxRegionX;
+        this.maxRegionZ = maxRegionZ;
+    }
+
+    static long regionKey(int regionX, int regionZ) {
+        return (((long) regionX) << 32) | (regionZ & 0xFFFFFFFFL);
+    }
+
+    /**
+     * Retains the dependency square of a FULL request at the chunk. The request loads every position in it
+     * anyway, so retaining loads nothing.
+     */
+    synchronized void retainAround(int chunkX, int chunkZ) {
+        for (int x = chunkX - radius; x <= chunkX + radius; x++) {
+            long[] bits = null;
+            long bitsKey = 0L;
+            for (int z = chunkZ - radius; z <= chunkZ + radius; z++) {
+                long key = regionKey(x >> 5, z >> 5);
+                if (bits == null || key != bitsKey) {
+                    bitsKey = key;
+                    bits = held.computeIfAbsent(key, ignored -> new long[REGION_WORDS]);
+                }
+                int index = ((z & 31) << 5) | (x & 31);
+                long mask = 1L << (index & 63);
+                if ((bits[index >> 6] & mask) != 0L) {
+                    continue;
+                }
+                bits[index >> 6] |= mask;
+                size++;
+                tickets.retain(x, z);
+            }
+        }
+
+        if (size > capacity) {
+            releaseFarthestDrained(chunkX >> 5, chunkZ >> 5);
+        }
+    }
+
+    /**
+     * Records a drained region and releases every retained position no undrained in-area region reaches.
+     */
+    synchronized void drained(int regionX, int regionZ) {
+        drained.add(regionKey(regionX, regionZ));
+        for (int x = regionX - 1; x <= regionX + 1; x++) {
+            for (int z = regionZ - 1; z <= regionZ + 1; z++) {
+                long key = regionKey(x, z);
+                long[] bits = held.get(key);
+                if (bits != null && release(key, bits, true)) {
+                    held.remove(key);
+                }
+            }
+        }
+    }
+
+    synchronized void releaseAll() {
+        for (Map.Entry<Long, long[]> entry : held.entrySet()) {
+            release(entry.getKey(), entry.getValue(), false);
+        }
+        held.clear();
+    }
+
+    synchronized int size() {
+        return size;
+    }
+
+    private void releaseFarthestDrained(int regionX, int regionZ) {
+        List<Long> candidates = new ArrayList<>();
+        for (Long key : held.keySet()) {
+            if (drained.contains(key)) {
+                candidates.add(key);
+            }
+        }
+        candidates.sort(Comparator.comparingLong((Long key) -> {
+            long dx = (int) (key >> 32) - regionX;
+            long dz = key.intValue() - regionZ;
+            return dx * dx + dz * dz;
+        }).reversed());
+        int target = capacity - capacity / 8;
+        for (Long key : candidates) {
+            if (size <= target) {
+                return;
+            }
+            release(key, held.remove(key), false);
+        }
+    }
+
+    private boolean release(long key, long[] bits, boolean onlyUnreachable) {
+        int baseX = ((int) (key >> 32)) << 5;
+        int baseZ = ((int) key) << 5;
+        boolean empty = true;
+        for (int word = 0; word < REGION_WORDS; word++) {
+            long remaining = bits[word];
+            while (remaining != 0L) {
+                int bit = Long.numberOfTrailingZeros(remaining);
+                remaining &= remaining - 1L;
+                int index = (word << 6) | bit;
+                int x = baseX + (index & 31);
+                int z = baseZ + (index >> 5);
+                if (onlyUnreachable && reachedByUndrainedRegion(x, z)) {
+                    empty = false;
+                    continue;
+                }
+                bits[word] &= ~(1L << bit);
+                size--;
+                tickets.release(x, z);
+            }
+        }
+        return empty;
+    }
+
+    private boolean reachedByUndrainedRegion(int chunkX, int chunkZ) {
+        for (int regionX = (chunkX - radius) >> 5; regionX <= (chunkX + radius) >> 5; regionX++) {
+            for (int regionZ = (chunkZ - radius) >> 5; regionZ <= (chunkZ + radius) >> 5; regionZ++) {
+                if (regionX >= minRegionX && regionX <= maxRegionX && regionZ >= minRegionZ && regionZ <= maxRegionZ
+                        && !drained.contains(regionKey(regionX, regionZ))) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
+    interface Tickets {
+        void retain(int chunkX, int chunkZ);
+
+        void release(int chunkX, int chunkZ);
+    }
+}

core/src/test/java/art/arcane/iris/world/pregen/AsyncPregenMethodConcurrencyCapTest.java

@@ -26,6 +26,14 @@ public class AsyncPregenMethodConcurrencyCapTest {
         assertEquals(256, AsyncPregenMethod.heapLimitedConcurrency(256, Long.MAX_VALUE));
     }
 
+    @Test
+    public void retainedChunksAreBudgetedAtOneMegabyteOfHeapEach() {
+        assertEquals(8192, AsyncPregenMethod.retentionCapacity(8L << 30));
+        assertEquals(16384, AsyncPregenMethod.retentionCapacity(16L << 30));
+        assertEquals(0, AsyncPregenMethod.retentionCapacity(0L));
+        assertEquals(Integer.MAX_VALUE, AsyncPregenMethod.retentionCapacity(Long.MAX_VALUE));
+    }
+
     @Test
     public void pregenWorkerTargetExpandsOnlyOptedInHighHeapJobs() {
         assertEquals(16, AsyncPregenMethod.pregenWorkerThreadTarget(16, (16L << 30) - 1L, true));

core/src/test/java/art/arcane/iris/world/pregen/AsyncPregenServerSaveOwnershipTest.java

@@ -17,7 +17,9 @@ import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.Queue;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentHashMap;
@@ -333,20 +335,39 @@ public class AsyncPregenServerSaveOwnershipTest {
     }
 
     @Test
-    public void completedRegionsUnloadWithoutWaitingForTheirNeighbors() throws Exception {
+    public void completedRegionsUnloadAtDrainWhileRetentionKeepsOnlyTheBordersUndrainedNeighboursNeed() throws Exception {
         Fixture fixture = new Fixture();
         fixture.chunks.clear();
+        Set<Long> retained = new HashSet<>();
+        PregenChunkRetention retention = new PregenChunkRetention(new PregenChunkRetention.Tickets() {
+            @Override
+            public void retain(int chunkX, int chunkZ) {
+                assertTrue(retained.add(PregenChunkRetention.regionKey(chunkX, chunkZ)));
+            }
+
+            @Override
+            public void release(int chunkX, int chunkZ) {
+                assertTrue(retained.remove(PregenChunkRetention.regionKey(chunkX, chunkZ)));
+            }
+        }, 10, 1_000_000, 0, 0, 127, 0);
+        fixture.set("retention", retention);
         try (SchedulingContext context = new SchedulingContext(fixture)) {
             for (int region = 0; region < 128; region++) {
                 long key = (long) region << 32;
+                retention.retainAround(region << 5, 0);
+                retention.retainAround((region << 5) + 31, 0);
                 fixture.pendingRegions.put(key, new AtomicInteger(2));
                 invoke(fixture.method, "onChunkCompleted", region << 5, 0, fixture.chunk);
                 assertEquals(1, fixture.pendingRegions.get(key).get());
                 fixture.method.onRegionSubmitted(region, 0);
                 assertTrue(fixture.pendingRegions.isEmpty());
                 assertTrue(fixture.chunks.isEmpty());
+                verify(fixture.binding, times(region + 1)).saveAndUnloadChunk(fixture.world, 2, 3);
+                for (long position : retained) {
+                    assertTrue((int) (position >> 32) >= ((region + 1) << 5) - 10);
                 }
-            verify(fixture.binding, times(128)).saveAndUnloadChunk(fixture.world, 2, 3);
+            }
+            assertTrue(retained.isEmpty());
         }
     }
 
@@ -391,6 +412,35 @@ public class AsyncPregenServerSaveOwnershipTest {
         }
     }
 
+    @Test
+    public void heapPressureReleasesEveryRetainedChunkBeforeEvicting() throws Exception {
+        Fixture fixture = new Fixture();
+        fixture.chunks.clear();
+        Set<Long> retained = new HashSet<>();
+        PregenChunkRetention retention = new PregenChunkRetention(new PregenChunkRetention.Tickets() {
+            @Override
+            public void retain(int chunkX, int chunkZ) {
+                retained.add(PregenChunkRetention.regionKey(chunkX, chunkZ));
+            }
+
+            @Override
+            public void release(int chunkX, int chunkZ) {
+                retained.remove(PregenChunkRetention.regionKey(chunkX, chunkZ));
+            }
+        }, 10, 1_000_000, 0, 0, 0, 0);
+        retention.retainAround(2, 3);
+        fixture.set("retention", retention);
+        fixture.pendingRegions.put(0L, new AtomicInteger(2));
+        try (SchedulingContext context = new SchedulingContext(fixture);
+             MockedStatic<MantleHeapPressure> pressure = mockStatic(MantleHeapPressure.class)) {
+            pressure.when(MantleHeapPressure::overHighWater).thenReturn(true);
+            invoke(fixture.method, "onChunkCompleted", 2, 3, fixture.chunk);
+            assertTrue(retained.isEmpty());
+            assertEquals(0, retention.size());
+            verify(fixture.binding).saveAndUnloadChunk(fixture.world, 2, 3);
+        }
+    }
+
     @Test
     public void heapPressureStartingAfterInitialCheckStillReclaimsCompletedChunks() throws Exception {
         Fixture fixture = new Fixture();

core/src/test/java/art/arcane/iris/world/pregen/PregenChunkRetentionTest.java

@@ -0,0 +1,142 @@
+package art.arcane.iris.world.pregen;
+
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class PregenChunkRetentionTest {
+    @Test
+    public void retainsEachPositionOfTheDependencySquareOnce() {
+        RecordingTickets tickets = new RecordingTickets();
+        PregenChunkRetention retention = new PregenChunkRetention(tickets, 2, 10_000, 0, 0, 0, 0);
+
+        retention.retainAround(5, 5);
+        assertEquals(25, tickets.retained.size());
+        assertEquals(square(3, 7, 3, 7), new HashSet<>(tickets.retained));
+
+        retention.retainAround(6, 5);
+        assertEquals(30, tickets.retained.size());
+        assertEquals(square(8, 8, 3, 7), new HashSet<>(tickets.retained.subList(25, 30)));
+        assertEquals(30, retention.size());
+    }
+
+    @Test
+    public void drainedRegionKeepsOnlyTheBorderAnUndrainedNeighbourStillReaches() {
+        RecordingTickets tickets = new RecordingTickets();
+        PregenChunkRetention retention = new PregenChunkRetention(tickets, 2, 10_000, 0, 0, 1, 0);
+        retention.retainAround(16, 16);
+        retention.retainAround(30, 16);
+
+        retention.drained(0, 0);
+        Set<Long> released = new HashSet<>(square(14, 18, 14, 18));
+        released.addAll(square(28, 29, 14, 18));
+        assertEquals(released, new HashSet<>(tickets.released));
+        assertEquals(15, retention.size());
+
+        retention.drained(1, 0);
+        assertEquals(new HashSet<>(tickets.retained), new HashSet<>(tickets.released));
+        assertEquals(tickets.retained.size(), tickets.released.size());
+        assertEquals(0, retention.size());
+    }
+
+    @Test
+    public void positionsOutsideTheAreaReleaseWithTheInAreaRegionsThatReachThem() {
+        RecordingTickets tickets = new RecordingTickets();
+        PregenChunkRetention retention = new PregenChunkRetention(tickets, 10, 10_000, 0, 0, 0, 0);
+        retention.retainAround(0, 0);
+        retention.retainAround(31, 31);
+
+        retention.drained(0, 0);
+        assertEquals(new HashSet<>(tickets.retained), new HashSet<>(tickets.released));
+        assertEquals(0, retention.size());
+    }
+
+    @Test
+    public void aLongRunOfRegionsOnlyHoldsTheActiveRegionAndTheStripBehindIt() {
+        RecordingTickets tickets = new RecordingTickets();
+        PregenChunkRetention retention = new PregenChunkRetention(tickets, 10, 1_000_000, 0, 0, 127, 0);
+        int peak = 0;
+        for (int region = 0; region < 128; region++) {
+            for (int x = 0; x < 32; x += 4) {
+                for (int z = 0; z < 32; z += 4) {
+                    retention.retainAround(region * 32 + x, z);
+                }
+            }
+            peak = Math.max(peak, retention.size());
+            retention.drained(region, 0);
+            int strip = region == 127 ? 0 : 17 * 49;
+            assertEquals(strip, retention.size());
+        }
+        assertEquals(49 * 49, peak);
+        assertEquals(new HashSet<>(tickets.retained), new HashSet<>(tickets.released));
+    }
+
+    @Test
+    public void overCapacityReleasesTheDrainedStripsFarthestFromTheActiveRegionFirst() {
+        RecordingTickets tickets = new RecordingTickets();
+        PregenChunkRetention retention = new PregenChunkRetention(tickets, 1, 20, 0, 0, 4, 0);
+        retention.retainAround(31, 5);
+        retention.drained(0, 0);
+        assertEquals(square(30, 30, 4, 6), new HashSet<>(tickets.released));
+        assertEquals(6, retention.size());
+
+        retention.retainAround(127, 5);
+        retention.drained(3, 0);
+        assertEquals(12, retention.size());
+
+        retention.retainAround(40, 5);
+        Set<Long> released = new HashSet<>(square(30, 30, 4, 6));
+        released.addAll(square(126, 127, 4, 6));
+        assertEquals(released, new HashSet<>(tickets.released));
+        assertEquals(18, retention.size());
+    }
+
+    @Test
+    public void releaseAllDropsEveryTicketAndRetainsAgainAfterwards() {
+        RecordingTickets tickets = new RecordingTickets();
+        PregenChunkRetention retention = new PregenChunkRetention(tickets, 1, 10_000, 0, 0, 0, 0);
+        retention.retainAround(4, 4);
+        retention.releaseAll();
+        assertEquals(new HashSet<>(tickets.retained), new HashSet<>(tickets.released));
+        assertEquals(0, retention.size());
+
+        retention.retainAround(4, 4);
+        assertEquals(18, tickets.retained.size());
+        assertEquals(9, retention.size());
+    }
+
+    private static Set<Long> square(int minX, int maxX, int minZ, int maxZ) {
+        Set<Long> positions = new HashSet<>();
+        for (int x = minX; x <= maxX; x++) {
+            for (int z = minZ; z <= maxZ; z++) {
+                positions.add(position(x, z));
+            }
+        }
+        return positions;
+    }
+
+    private static long position(int x, int z) {
+        return PregenChunkRetention.regionKey(x, z);
+    }
+
+    private static final class RecordingTickets implements PregenChunkRetention.Tickets {
+        private final List<Long> retained = new ArrayList<>();
+        private final List<Long> released = new ArrayList<>();
+
+        @Override
+        public void retain(int chunkX, int chunkZ) {
+            retained.add(position(chunkX, chunkZ));
+        }
+
+        @Override
+        public void release(int chunkX, int chunkZ) {
+            released.add(position(chunkX, chunkZ));
+        }
+    }
+}