Fixed hydrology shutdown

- Hydrology close drains active tile plans and column composition before returning
- Queued tile plans cancel without waiting for their executor to start them
- Closed caches reject new tile and column demand
- Shutdown tests cover background plans, queued demand, active composition, and recursive close
This commit is contained in:
Brian Neumann-Fopiano
2026-09-06 03:15:15 -04:00
parent 17c927acdd
commit dde3c149ad
2 changed files with 177 additions and 9 deletions
@@ -16,6 +16,7 @@ import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;
@@ -131,6 +132,7 @@ public final class HydrologyTileCache implements AutoCloseable {
*/
public HydrologyTile get(HydrologyTileKey key) {
Objects.requireNonNull(key, "key");
requireOpen();
if (prefetchExecutor == null) {
return planInline(key, cacheEpoch.get());
}
@@ -156,7 +158,11 @@ public final class HydrologyTileCache implements AutoCloseable {
}
CacheLoadKey<HydrologyTileKey> loadKey = new CacheLoadKey<>(epoch, key);
PendingLoad<HydrologyTile> owned = new PendingLoad<>(Thread.currentThread(), new CompletableFuture<>());
PendingLoad<HydrologyTile> existing = loading.putIfAbsent(loadKey, owned);
PendingLoad<HydrologyTile> existing;
synchronized (publicationLock) {
requireOpen();
existing = loading.putIfAbsent(loadKey, owned);
}
if (existing != null) {
return awaitLoad(existing);
}
@@ -376,10 +382,45 @@ public final class HydrologyTileCache implements AutoCloseable {
@Override
public void close() {
closed.set(true);
ArrayList<CompletableFuture<?>> active = new ArrayList<>();
ArrayList<CompletableFuture<HydrologyTile>> queued = new ArrayList<>();
synchronized (publicationLock) {
for (PendingLoad<HydrologyTile> load : loading.values()) {
collectActiveLoad(active, load);
}
for (PendingLoad<ChunkColumns> load : composing.values()) {
collectActiveLoad(active, load);
}
closed.set(true);
planning.forEach((key, future) -> {
if (!loading.containsKey(key)) {
queued.add(future);
}
});
}
for (CompletableFuture<HydrologyTile> future : queued) {
future.completeExceptionally(new CancellationException("Hydrology tile cache is closed."));
}
CompletableFuture.allOf(active.toArray(CompletableFuture<?>[]::new))
.handle((ignored, failure) -> null)
.join();
planning.clear();
clear();
}
private static void collectActiveLoad(List<CompletableFuture<?>> active, PendingLoad<?> load) {
if (load.owner() == Thread.currentThread()) {
throw new IllegalStateException("Hydrology tile cache cannot close from an active load.");
}
active.add(load.future());
}
private void requireOpen() {
if (closed.get()) {
throw new CancellationException("Hydrology tile cache is closed.");
}
}
private boolean validSharedTile(HydrologyTile tile, SharedTileKey sharedKey) {
SharedCacheScope scope = sharedKey.scope();
return tile.key().equals(sharedKey.tileKey())
@@ -389,6 +430,7 @@ public final class HydrologyTileCache implements AutoCloseable {
}
private ChunkColumns chunkColumns(int blockX, int blockZ) {
requireOpen();
int chunkX = Math.floorDiv(blockX, CHUNK_SIZE);
int chunkZ = Math.floorDiv(blockZ, CHUNK_SIZE);
long epoch = cacheEpoch.get();
@@ -420,7 +462,11 @@ public final class HydrologyTileCache implements AutoCloseable {
return present;
}
PendingLoad<ChunkColumns> owned = new PendingLoad<>(Thread.currentThread(), new CompletableFuture<>());
PendingLoad<ChunkColumns> existing = composing.putIfAbsent(loadKey, owned);
PendingLoad<ChunkColumns> existing;
synchronized (publicationLock) {
requireOpen();
existing = composing.putIfAbsent(loadKey, owned);
}
if (existing != null) {
return awaitLoad(existing);
}
@@ -652,6 +698,7 @@ public final class HydrologyTileCache implements AutoCloseable {
*/
public List<HydrologyTile> tiles(List<HydrologyTileKey> keys) {
Objects.requireNonNull(keys, "keys");
requireOpen();
HydrologyTile[] loaded = new HydrologyTile[keys.size()];
boolean missing = false;
for (int index = 0; index < loaded.length; index++) {
@@ -724,17 +771,29 @@ public final class HydrologyTileCache implements AutoCloseable {
return CompletableFuture.completedFuture(present);
}
CompletableFuture<HydrologyTile> future = new CompletableFuture<>();
CompletableFuture<HydrologyTile> existing = planning.putIfAbsent(loadKey, future);
CompletableFuture<HydrologyTile> existing;
synchronized (publicationLock) {
if (closed.get()) {
return CompletableFuture.failedFuture(new CancellationException("Hydrology tile cache is closed."));
}
existing = planning.putIfAbsent(loadKey, future);
}
if (existing != null) {
return existing;
}
try {
prefetchExecutor.execute(() -> {
if (future.isDone()) {
planning.remove(loadKey, future);
return;
}
try {
future.complete(planInline(key, loadKey.epoch()));
} catch (Throwable failure) {
future.completeExceptionally(failure);
IrisLogging.reportError(failure);
if (!(failure instanceof CancellationException) || !closed.get()) {
IrisLogging.reportError(failure);
}
} finally {
planning.remove(loadKey, future);
}
@@ -15,6 +15,7 @@ import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ExecutionException;
@@ -23,6 +24,7 @@ import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
@@ -43,6 +45,104 @@ import static org.mockito.Mockito.never;
import static org.mockito.Mockito.when;
public class HydrologyTileCacheTest {
@Test
public void closeDrainsBackgroundPlanningAndRejectsNewDemand() throws Exception {
HydrologyPlanner planner = mock(HydrologyPlanner.class);
HydrologyTile tile = mock(HydrologyTile.class);
HydrologyTileKey key = new HydrologyTileKey(0, 0);
CountDownLatch started = new CountDownLatch(1);
CountDownLatch release = new CountDownLatch(1);
when(planner.plan(key)).thenAnswer(invocation -> {
started.countDown();
assertTrue(release.await(5, TimeUnit.SECONDS));
return tile;
});
ExecutorService workers = Executors.newSingleThreadExecutor();
ExecutorService callers = Executors.newFixedThreadPool(2);
HydrologyTileCache cache = new HydrologyTileCache(planner, 4, workers);
try {
Future<HydrologyTile> result = callers.submit(() -> cache.get(key));
assertTrue(started.await(5, TimeUnit.SECONDS));
Future<?> closing = callers.submit(cache::close);
awaitClosed(cache);
assertThrows(TimeoutException.class, () -> closing.get(100, TimeUnit.MILLISECONDS));
assertThrows(CancellationException.class, () -> cache.get(new HydrologyTileKey(1, 0)));
assertThrows(CancellationException.class, () -> cache.columnAt(0, 0));
release.countDown();
assertSame(tile, result.get(5, TimeUnit.SECONDS));
closing.get(5, TimeUnit.SECONDS);
assertEquals(0, cache.size());
verify(planner, never()).plan(new HydrologyTileKey(1, 0));
} finally {
release.countDown();
workers.shutdownNow();
callers.shutdownNow();
assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS));
assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS));
}
}
@Test
public void closeWaitsForColumnCompositionBeforeReturning() throws Exception {
HydrologyPlanner planner = mock(HydrologyPlanner.class);
HydrologyTile tile = mock(HydrologyTile.class);
when(planner.settings()).thenReturn(emptySettings());
when(planner.plan(any(HydrologyTileKey.class))).thenReturn(tile);
CountDownLatch started = new CountDownLatch(1);
CountDownLatch release = new CountDownLatch(1);
when(tile.columnAt(anyInt(), anyInt())).thenAnswer(invocation -> {
if (invocation.<Integer>getArgument(0) == 0 && invocation.<Integer>getArgument(1) == 0) {
started.countDown();
assertTrue(release.await(5, TimeUnit.SECONDS));
}
return Optional.empty();
});
HydrologyTileCache cache = new HydrologyTileCache(planner, 4);
ExecutorService callers = Executors.newFixedThreadPool(2);
try {
Future<Optional<HydrologyColumnSample>> result = callers.submit(() -> cache.columnAt(0, 0));
assertTrue(started.await(5, TimeUnit.SECONDS));
Future<?> closing = callers.submit(cache::close);
awaitClosed(cache);
assertThrows(TimeoutException.class, () -> closing.get(100, TimeUnit.MILLISECONDS));
release.countDown();
assertTrue(result.get(5, TimeUnit.SECONDS).isEmpty());
closing.get(5, TimeUnit.SECONDS);
assertEquals(0, cache.size());
} finally {
release.countDown();
callers.shutdownNow();
assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS));
}
}
@Test
public void closingFromAnOwnedLoadFailsWithoutDeadlockingOrClosingTheCache() {
HydrologyPlanner planner = mock(HydrologyPlanner.class);
HydrologyTile tile = mock(HydrologyTile.class);
HydrologyTileCache cache = new HydrologyTileCache(planner, 4);
HydrologyTileKey key = new HydrologyTileKey(0, 0);
when(planner.plan(key)).thenAnswer(invocation -> {
assertThrows(IllegalStateException.class, cache::close);
return tile;
});
assertSame(tile, cache.get(key));
assertSame(tile, cache.get(key));
cache.close();
assertThrows(CancellationException.class, () -> cache.get(key));
}
private static void awaitClosed(HydrologyTileCache cache) throws Exception {
Field field = HydrologyTileCache.class.getDeclaredField("closed");
field.setAccessible(true);
AtomicBoolean closed = (AtomicBoolean) field.get(cache);
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
while (!closed.get() && System.nanoTime() < deadline) {
Thread.sleep(1);
}
assertTrue(closed.get());
}
@Test
public void explicitAreaIsFullyQueuedBeforeExecutorCallbacksCanAddOldWork() throws Exception {
for (boolean pregeneration : new boolean[]{false, true}) {
@@ -265,7 +365,7 @@ public class HydrologyTileCacheTest {
}
@Test
public void planCompletingAfterCloseDoesNotEnterTheSharedStudioCache() throws Exception {
public void closeWaitsForActivePlansWithoutPublishingTheirSharedTiles() throws Exception {
HydrologyPlanner closingPlanner = mock(HydrologyPlanner.class);
HydrologyPlanner succeedingPlanner = mock(HydrologyPlanner.class);
HydrologyTile lateTile = mock(HydrologyTile.class);
@@ -287,13 +387,15 @@ public class HydrologyTileCacheTest {
when(lateTile.tileSize()).thenReturn(64);
HydrologyTileCache closing = new HydrologyTileCache(closingPlanner, 4, null, null, scope);
HydrologyTileCache succeeding = new HydrologyTileCache(succeedingPlanner, 4, null, null, scope);
ExecutorService executor = Executors.newSingleThreadExecutor();
ExecutorService executor = Executors.newFixedThreadPool(2);
try {
Future<HydrologyTile> late = executor.submit(() -> closing.get(key));
assertTrue(planningStarted.await(10L, TimeUnit.SECONDS));
closing.close();
Future<?> shutdown = executor.submit(closing::close);
assertThrows(TimeoutException.class, () -> shutdown.get(100L, TimeUnit.MILLISECONDS));
allowCompletion.countDown();
assertSame(lateTile, late.get(10L, TimeUnit.SECONDS));
shutdown.get(10L, TimeUnit.SECONDS);
assertEquals(0, closing.size());
assertSame(plannedTile, succeeding.get(key));
verify(succeedingPlanner, times(1)).plan(key);
@@ -1451,7 +1553,14 @@ public class HydrologyTileCacheTest {
freshPrefetch.run();
}
oldDemand.run();
assertSame(tile, result.get(5, TimeUnit.SECONDS));
if (close) {
ExecutionException failure = assertThrows(ExecutionException.class,
() -> result.get(5, TimeUnit.SECONDS));
assertTrue(failure.getCause() instanceof CancellationException);
verify(planner, never()).plan(new HydrologyTileKey(1, 0));
} else {
assertSame(tile, result.get(5, TimeUnit.SECONDS));
}
assertTrue(queued.isEmpty());
assertEquals(close ? 0 : 1, cache.size());
verify(planner, never()).plan(new HydrologyTileKey(2, 0));