Retained failed generation ownership

- Keep writers until cancelled components actually finish
- Retain incomplete runtime and target ownership for retries
- Gate storage release on completed service drains and saves

core/src/main/java/art/arcane/iris/engine/EngineHotloader.java

@@ -175,7 +175,7 @@ final class EngineHotloader {
                 broadcastStudioHotload(false, "");
             } catch (Throwable e) {
                 if (!published) {
-                    if (replacementData != null) {
+                    if (replacementData != null && !engine.shutdownSequence.retainsData(replacementData)) {
                         replacementData.unregisterEngine(engine);
                         Throwable replacementDataFailure = runCleanup(null, replacementData::close);
                         if (replacementDataFailure != null) {

core/src/main/java/art/arcane/iris/engine/EngineRuntimeBuilder.java

@@ -209,7 +209,12 @@ final class EngineRuntimeBuilder {
                 next.generation().mantle(),
                 null);
         if (retirementFailure != null) {
-            retirementFailure = engine.shutdownSequence.closeRuntime(next, retirementFailure);
+            Throwable cleanupFailure = engine.shutdownSequence.closeRuntime(
+                    next, previous == null ? null : previous.generation().mantle(), null);
+            if (cleanupFailure != null) {
+                engine.shutdownSequence.retainUnpublishedRuntime(next);
+                retirementFailure = EngineShutdownSequence.appendFailure(retirementFailure, cleanupFailure);
+            }
             engine.lifecycleState = LifecycleState.FAILED;
             throw new IllegalStateException("Failed to retire the previous Iris engine runtime.", retirementFailure);
         }
@@ -229,6 +234,7 @@ final class EngineRuntimeBuilder {
             engine.getClosing().set(true);
             engine.backgroundTasks.closeBackgroundTaskAdmission();
             engine.lifecycleState = LifecycleState.FAILED;
+            Throwable cleanupFailure = null;
             try {
                 engine.getGenerationSessions().sealAndAwait(
                         "failed world manager start",
@@ -236,11 +242,17 @@ final class EngineRuntimeBuilder {
                         true
                 );
             } catch (Throwable drainFailure) {
-                e.addSuppressed(drainFailure);
+                cleanupFailure = drainFailure;
             }
-            engine.shutdownSequence.closeRuntime(next, e);
-            if (engine.runtime == next) {
-                engine.runtime = null;
+            if (cleanupFailure == null) {
+                cleanupFailure = engine.shutdownSequence.closeRuntime(next, null);
+            }
+            if (cleanupFailure == null) {
+                if (engine.runtime == next) {
+                    engine.runtime = null;
+                }
+            } else {
+                e.addSuppressed(cleanupFailure);
             }
             throw new IllegalStateException("Failed to start the Iris world manager.", e);
         }

core/src/main/java/art/arcane/iris/engine/EngineShutdownSequence.java

@@ -18,6 +18,7 @@
 
 package art.arcane.iris.engine;
 
+import art.arcane.iris.core.loader.IrisData;
 import art.arcane.iris.engine.EngineBackgroundTasks.BackgroundTaskDrain;
 import art.arcane.iris.engine.EngineRuntimeBuilder.RuntimeAssembly;
 import art.arcane.iris.engine.IrisEngine.LifecycleState;
@@ -29,6 +30,11 @@ import art.arcane.iris.engine.mantle.EngineMantle;
 import art.arcane.iris.spi.IrisLogging;
 import art.arcane.iris.spi.IrisServices;
 
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+
 /**
  * Owns the ordered teardown of a single {@link IrisEngine}, both for a normal close and for a
  * construction that failed part way through. Every resource release is gated on the previous one
@@ -38,6 +44,10 @@ final class EngineShutdownSequence {
     private static final long CLOSE_RETRY_DRAIN_TIMEOUT_MILLIS = 5000L;
 
     private final IrisEngine engine;
+    private final Set<RuntimeAssembly> incompleteAssemblies = ConcurrentHashMap.newKeySet();
+    private final Set<EngineRuntime> unpublishedRuntimes = ConcurrentHashMap.newKeySet();
+    private final Set<EngineTarget> incompleteTargets = Collections.synchronizedSet(
+            Collections.newSetFromMap(new IdentityHashMap<>()));
     private boolean runtimeReleased;
     private boolean targetReleased;
     private boolean engineDataReleased;
@@ -172,17 +182,21 @@ final class EngineShutdownSequence {
         EngineTickRegistry.unregisterTicking(engine);
         engine.lifecycleState = LifecycleState.FAILED;
         Throwable cleanupFailure = null;
+        boolean sessionsDrained = true;
         try {
             engine.getGenerationSessions().sealAndAwait("failed initialization", 0L, true);
         } catch (Throwable e) {
+            sessionsDrained = false;
             cleanupFailure = appendFailure(cleanupFailure, e);
         }
         engine.backgroundTasks.cancelBackgroundTasks("failed initialization");
         BackgroundTaskDrain backgroundDrain = engine.backgroundTasks.drainBackgroundTasks("failed initialization");
         cleanupFailure = appendFailure(cleanupFailure, backgroundDrain.failure());
-        if (!backgroundDrain.allowsResourceRelease()) {
-            cleanupFailure = appendFailure(cleanupFailure,
-                    new IllegalStateException("Iris background tasks remain active after failed initialization."));
+        if (!sessionsDrained || !backgroundDrain.allowsResourceRelease()) {
+            if (!backgroundDrain.allowsResourceRelease()) {
+                cleanupFailure = appendFailure(cleanupFailure,
+                        new IllegalStateException("Iris background tasks remain active after failed initialization."));
+            }
             if (cleanupFailure != original) {
                 original.addSuppressed(cleanupFailure);
             }
@@ -192,27 +206,19 @@ final class EngineShutdownSequence {
                 () -> NativeStructureOwnershipStore.close(engine));
         cleanupFailure = appendFailure(cleanupFailure, ownershipFailure);
         if (ownershipFailure == null) {
-            Throwable runtimeFailure = closeDetachedGenerationRuntimes(null);
-            if (runtimeFailure == null) {
-                runtimeFailure = closeRuntime(engine.runtime, null);
+            cleanupFailure = releaseRuntime(cleanupFailure);
+            if (runtimeReleased) {
+                cleanupFailure = releaseTarget(cleanupFailure);
             }
-            cleanupFailure = appendFailure(cleanupFailure, runtimeFailure);
-            if (runtimeFailure != null) {
-                if (cleanupFailure != original) {
-                    original.addSuppressed(cleanupFailure);
-                }
-                return;
+            if (targetReleased) {
+                cleanupFailure = releaseEngineDataForShutdown(cleanupFailure);
+            }
+            if (engineDataReleased) {
+                cleanupFailure = releasePreservation(cleanupFailure);
+            }
+            if (runtimeReleased && targetReleased && engineDataReleased && preservationReleased) {
+                engine.closed = true;
             }
-            engine.runtime = null;
-            cleanupFailure = runCleanup(cleanupFailure, engine.publishedTarget::close);
-            cleanupFailure = runCleanup(cleanupFailure, engine.engineDataStore::releaseEngineData);
-            engine.closed = true;
-            cleanupFailure = runCleanup(cleanupFailure, () -> {
-                PreservationRegistry registry = IrisServices.getOrNull(PreservationRegistry.class);
-                if (registry != null) {
-                    registry.dereference();
-                }
-            });
         }
         if (cleanupFailure != null && cleanupFailure != original) {
             original.addSuppressed(cleanupFailure);
@@ -223,41 +229,55 @@ final class EngineShutdownSequence {
         if (assembly == null) {
             return failure;
         }
-        failure = runCleanup(failure, () -> {
-            if (assembly.worldManager != null) {
-                assembly.worldManager.close();
-            }
-        });
-        failure = runCleanup(failure, () -> {
-            if (assembly.effects != null) {
-                assembly.effects.close();
-            }
-        });
-        failure = runCleanup(failure, () -> {
-            if (assembly.mode != null) {
-                assembly.mode.close();
-            }
-        });
-        Throwable complexFailure = runCleanup(null, () -> {
-            if (assembly.complex != null) {
-                assembly.complex.close();
-            }
-        });
-        failure = appendFailure(failure, complexFailure);
-        if (complexFailure != null) {
-            return failure;
+        incompleteAssemblies.add(assembly);
+        RuntimeAssembly previous = engine == null ? null : engine.runtimeAssembly.get();
+        if (engine != null) {
+            engine.runtimeAssembly.set(assembly);
         }
-        failure = runCleanup(failure, () -> {
-            if (assembly.hash32 != null) {
-                assembly.hash32.cancel(true);
+        Throwable assemblyFailure;
+        try {
+            assemblyFailure = closeAssemblyResources(assembly);
+        } finally {
+            if (engine != null) {
+                if (previous == null) {
+                    engine.runtimeAssembly.remove();
+                } else {
+                    engine.runtimeAssembly.set(previous);
+                }
             }
-        });
-        if (assembly.mantle != null && assembly.ownsMantle) {
-            assembly.ownsMantle = false;
-            failure = runCleanup(failure, assembly.mantle::saveAllNow);
-            failure = runCleanup(failure, assembly.mantle::close);
         }
-        return failure;
+        if (assemblyFailure == null) {
+            incompleteAssemblies.remove(assembly);
+        } else if (engine != null) {
+            engine.lifecycleState = LifecycleState.FAILED;
+        }
+        return appendFailure(failure, assemblyFailure);
+    }
+
+    void retainUnpublishedRuntime(EngineRuntime runtime) {
+        unpublishedRuntimes.add(runtime);
+    }
+
+    boolean retainsData(IrisData data) {
+        if (engine.runtime != null && engine.runtime.generation().data() == data) {
+            return true;
+        }
+        for (RuntimeAssembly assembly : incompleteAssemblies) {
+            if (assembly.target.getData() == data) {
+                return true;
+            }
+        }
+        for (EngineRuntime runtime : unpublishedRuntimes) {
+            if (runtime.generation().data() == data) {
+                return true;
+            }
+        }
+        for (EngineTarget target : incompleteTargets.toArray(new EngineTarget[0])) {
+            if (target.getData() == data) {
+                return true;
+            }
+        }
+        return false;
     }
 
     Throwable closeRuntime(EngineRuntime engineRuntime, Throwable failure) {
@@ -272,8 +292,11 @@ final class EngineShutdownSequence {
         if (engineRuntime == null) {
             return failure;
         }
-        failure = runCleanup(failure, engineRuntime.worldManager()::close);
-        failure = runCleanup(failure, engineRuntime.effects()::close);
+        Throwable serviceFailure = runCleanup(null, engineRuntime.worldManager()::close);
+        serviceFailure = runCleanup(serviceFailure, engineRuntime.effects()::close);
+        if (serviceFailure != null) {
+            return appendFailure(failure, serviceFailure);
+        }
         return closeGenerationRuntime(engineRuntime.generation(), retainedMantle, failure);
     }
 
@@ -289,7 +312,10 @@ final class EngineShutdownSequence {
         if (generationRuntime == null) {
             return failure;
         }
-        failure = runCleanup(failure, generationRuntime.mode()::close);
+        Throwable modeFailure = runCleanup(null, generationRuntime.mode()::close);
+        if (modeFailure != null) {
+            return appendFailure(failure, modeFailure);
+        }
         Throwable complexFailure = runCleanup(null, generationRuntime.complex()::close);
         failure = appendFailure(failure, complexFailure);
         if (complexFailure != null) {
@@ -297,8 +323,7 @@ final class EngineShutdownSequence {
         }
         failure = runCleanup(failure, () -> generationRuntime.hash32().cancel(true));
         if (generationRuntime.mantle() != retainedMantle) {
-            failure = runCleanup(failure, generationRuntime.mantle()::saveAllNow);
-            failure = runCleanup(failure, generationRuntime.mantle()::close);
+            failure = appendFailure(failure, closeMantle(generationRuntime.mantle()));
         }
         return failure;
     }
@@ -320,6 +345,18 @@ final class EngineShutdownSequence {
     }
 
     Throwable closeDetachedTarget(EngineTarget target, Throwable failure) {
+        for (RuntimeAssembly assembly : incompleteAssemblies) {
+            if (assembly.target == target) {
+                return appendFailure(failure,
+                        new IllegalStateException("Iris runtime assembly still owns the detached target."));
+            }
+        }
+        for (EngineRuntime runtime : unpublishedRuntimes) {
+            if (runtime.generation().target() == target) {
+                return appendFailure(failure,
+                        new IllegalStateException("Iris unpublished runtime still owns the detached target."));
+            }
+        }
         failure = runCleanup(failure, () -> target.getData().unregisterEngine(engine));
         failure = runCleanup(failure, target::close);
         failure = runCleanup(failure, target.getData()::close);
@@ -330,6 +367,10 @@ final class EngineShutdownSequence {
         if (runtimeReleased) {
             return failure;
         }
+        Throwable incompleteFailure = closeIncompleteRuntimes();
+        if (incompleteFailure != null) {
+            return appendFailure(failure, incompleteFailure);
+        }
         Throwable detachedFailure = closeDetachedGenerationRuntimes(null);
         if (detachedFailure != null) {
             return appendFailure(failure, detachedFailure);
@@ -343,6 +384,86 @@ final class EngineShutdownSequence {
         return failure;
     }
 
+    private Throwable closeAssemblyResources(RuntimeAssembly assembly) {
+        Throwable failure = runCleanup(null, () -> {
+            if (assembly.worldManager != null) {
+                assembly.worldManager.close();
+            }
+        });
+        failure = runCleanup(failure, () -> {
+            if (assembly.effects != null) {
+                assembly.effects.close();
+            }
+        });
+        if (failure != null) {
+            return failure;
+        }
+        failure = runCleanup(null, () -> {
+            if (assembly.mode != null) {
+                assembly.mode.close();
+            }
+        });
+        if (failure != null) {
+            return failure;
+        }
+        failure = runCleanup(null, () -> {
+            if (assembly.complex != null) {
+                assembly.complex.close();
+            }
+        });
+        if (failure != null) {
+            return failure;
+        }
+        failure = runCleanup(null, () -> {
+            if (assembly.hash32 != null) {
+                assembly.hash32.cancel(true);
+            }
+        });
+        if (assembly.mantle != null && assembly.ownsMantle) {
+            Throwable mantleFailure = closeMantle(assembly.mantle);
+            if (mantleFailure == null) {
+                assembly.ownsMantle = false;
+            }
+            failure = appendFailure(failure, mantleFailure);
+        }
+        return failure;
+    }
+
+    private Throwable closeIncompleteRuntimes() {
+        Throwable failure = null;
+        for (RuntimeAssembly assembly : incompleteAssemblies.toArray(new RuntimeAssembly[0])) {
+            Throwable assemblyFailure = closeAssembly(assembly, null);
+            if (assemblyFailure == null && assembly.target != engine.publishedTarget) {
+                incompleteTargets.add(assembly.target);
+            }
+            failure = appendFailure(failure, assemblyFailure);
+        }
+        EngineMantle retainedMantle = engine.runtime == null ? null : engine.runtime.generation().mantle();
+        for (EngineRuntime unpublished : unpublishedRuntimes.toArray(new EngineRuntime[0])) {
+            Throwable runtimeFailure = closeRuntime(unpublished, retainedMantle, null);
+            if (runtimeFailure == null) {
+                unpublishedRuntimes.remove(unpublished);
+                if (unpublished.generation().target() != engine.publishedTarget) {
+                    incompleteTargets.add(unpublished.generation().target());
+                }
+            }
+            failure = appendFailure(failure, runtimeFailure);
+        }
+        for (EngineTarget target : incompleteTargets.toArray(new EngineTarget[0])) {
+            Throwable targetFailure = closeDetachedTarget(target, null);
+            if (targetFailure == null) {
+                incompleteTargets.remove(target);
+            }
+            failure = appendFailure(failure, targetFailure);
+        }
+        return failure;
+    }
+
+    private static Throwable closeMantle(EngineMantle mantle) {
+        Throwable failure = runCleanup(null, mantle::saveAllNow);
+        return failure == null ? runCleanup(null, mantle::close) : failure;
+    }
+
     private Throwable closeDetachedGenerationRuntimes(Throwable failure) {
         GenerationRuntime[] detachedRuntimes = engine.detachedGenerationRuntimes.toArray(new GenerationRuntime[0]);
         for (GenerationRuntime detached : detachedRuntimes) {

core/src/main/java/art/arcane/iris/engine/mantle/MatterGenerator.java

@@ -22,12 +22,13 @@ import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.Future;
+import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
 public interface MatterGenerator {
     MultiBurst DISPATCHER = MultiBurst.burst;
-    ConcurrentHashMap<MatterTaskKey, CompletableFuture<Void>> IN_FLIGHT_COMPONENTS = new ConcurrentHashMap<>();
+    ConcurrentHashMap<MatterTaskKey, MatterComponentTask> IN_FLIGHT_COMPONENTS = new ConcurrentHashMap<>();
     long COMPONENT_TASK_POLL_MS = 1000L;
     long COMPONENT_TASK_TIMEOUT_MS = Long.getLong("iris.mantle.componentTimeout", 120000L);
 
@@ -368,38 +369,37 @@ public interface MatterGenerator {
     ) {
         MantleFlag flag = component.getFlag();
         MatterTaskKey key = new MatterTaskKey(getMantle(), chunkX, chunkZ, flag);
+        MatterComponentTask task = new MatterComponentTask(key);
         if (chunk.isFlagged(flag)) {
-            return new MatterComponentTask(key, CompletableFuture.completedFuture(null), null);
+            task.finish(null);
+            return task;
         }
 
-        CompletableFuture<Void> future = new CompletableFuture<>();
-        CompletableFuture<Void> existing = IN_FLIGHT_COMPONENTS.putIfAbsent(key, future);
+        MatterComponentTask existing = IN_FLIGHT_COMPONENTS.putIfAbsent(key, task);
         if (existing != null) {
-            return new MatterComponentTask(key, existing, null);
+            return existing;
         }
 
         try {
             if (DISPATCHER.ownsCurrentThread()) {
                 // A pool thread runs its claim inline: a task queued behind pool threads that block
                 // on unmanaged waits (structure builds, cache locks) could starve the whole pool.
-                completeComponentTask(future, key, chunk, component, writer, chunkX, chunkZ, context, IrisContext.get());
-                return new MatterComponentTask(key, future, null);
+                completeComponentTask(task, chunk, component, writer, chunkX, chunkZ, context, IrisContext.get());
+                return task;
             }
             IrisContext callerContext = IrisContext.get();
-            Future<?> submission = DISPATCHER.submit(() -> completeComponentTask(
-                    future,
-                    key,
+            task.setSubmission(DISPATCHER.submit(() -> completeComponentTask(
+                    task,
                     chunk,
                     component,
                     writer,
                     chunkX,
                     chunkZ,
                     context,
-                    callerContext));
-            return new MatterComponentTask(key, future, submission);
+                    callerContext)));
+            return task;
         } catch (Throwable throwable) {
-            IN_FLIGHT_COMPONENTS.remove(key, future);
-            future.completeExceptionally(throwable);
+            task.cancel(throwable);
             throw throwable;
         }
     }
@@ -449,7 +449,7 @@ public interface MatterGenerator {
     }
 
     private static void awaitComponentTask(MatterComponentTask task) {
-        CompletableFuture<Void> future = task.future();
+        CompletableFuture<Void> future = task.future;
         if (future.isDone() && !future.isCompletedExceptionally()) {
             return;
         }
@@ -459,11 +459,17 @@ public interface MatterGenerator {
             try {
                 ForkJoinPool.managedBlock(new ComponentTaskBlocker(task));
             } catch (InterruptedException interruption) {
-                Thread.currentThread().interrupt();
+                try {
+                    task.await(COMPONENT_TASK_TIMEOUT_MS);
+                } finally {
+                    Thread.currentThread().interrupt();
+                }
+            } catch (RejectedExecutionException exhaustedPool) {
+                task.await(COMPONENT_TASK_TIMEOUT_MS);
             }
             return;
         }
-        awaitComponentTaskBlocking(task);
+        task.await(COMPONENT_TASK_TIMEOUT_MS);
     }
 
     /**
@@ -478,70 +484,22 @@ public interface MatterGenerator {
 
         @Override
         public boolean block() {
-            awaitComponentTaskBlocking(task);
+            task.await(COMPONENT_TASK_TIMEOUT_MS);
             return true;
         }
 
         @Override
         public boolean isReleasable() {
-            if (!task.future().isDone()) {
+            if (!task.future.isDone()) {
                 return false;
             }
-            task.future().join();
+            task.future.join();
             return true;
         }
     }
 
-    private static void awaitComponentTaskBlocking(MatterComponentTask task) {
-        CompletableFuture<Void> future = task.future();
-        long start = System.currentTimeMillis();
-        boolean interrupted = false;
-
-        try {
-            while (true) {
-                try {
-                    future.get(COMPONENT_TASK_POLL_MS, TimeUnit.MILLISECONDS);
-                    return;
-                } catch (InterruptedException interruption) {
-                    // The writer cannot close while this task still writes through it, so the wait is
-                    // uninterruptible and the flag is restored on the way out.
-                    interrupted = true;
-                } catch (TimeoutException timeout) {
-                    Future<?> submission = task.submission();
-                    boolean dropped = submission != null && submission.isDone();
-                    if (future.isDone()) {
-                        continue;
-                    }
-
-                    long waited = System.currentTimeMillis() - start;
-                    if (!dropped && waited < COMPONENT_TASK_TIMEOUT_MS) {
-                        continue;
-                    }
-
-                    // The task can no longer complete (the dispatcher dropped it) or is wedged. Drop
-                    // the dedup entry so later generations do not inherit a future that never
-                    // completes, and release anything else already waiting on it.
-                    IllegalStateException failure = new IllegalStateException("Mantle component " + task.key()
-                            + (dropped ? " was dropped by the dispatcher" : " did not complete in " + waited + "ms"));
-                    IN_FLIGHT_COMPONENTS.remove(task.key(), future);
-                    future.completeExceptionally(failure);
-                    IrisLogging.error(failure.getMessage());
-                    throw new CompletionException(failure);
-                } catch (ExecutionException executionFailure) {
-                    Throwable cause = executionFailure.getCause();
-                    throw new CompletionException(cause != null ? cause : executionFailure);
-                }
-            }
-        } finally {
-            if (interrupted) {
-                Thread.currentThread().interrupt();
-            }
-        }
-    }
-
     private void completeComponentTask(
-            CompletableFuture<Void> future,
-            MatterTaskKey key,
+            MatterComponentTask task,
             MantleChunk<Matter> chunk,
             MantleComponent component,
             MantleWriter writer,
@@ -550,14 +508,17 @@ public interface MatterGenerator {
             ChunkContext context,
             IrisContext callerContext
     ) {
+        if (!task.start()) {
+            return;
+        }
+        Throwable failure = null;
         try {
             runComponentWithContext(chunk, component, writer, chunkX, chunkZ, context, callerContext);
-            future.complete(null);
         } catch (Throwable throwable) {
-            future.completeExceptionally(throwable);
+            failure = throwable;
             throw throwable;
         } finally {
-            IN_FLIGHT_COMPONENTS.remove(key, future);
+            task.finish(failure);
         }
     }
 
@@ -597,11 +558,98 @@ public interface MatterGenerator {
         ));
     }
 
-    /**
-     * A launched component task. {@code submission} is the dispatcher handle when this generation
-     * owns the task, and null when the future belongs to another generation or already completed.
-     */
-    record MatterComponentTask(MatterTaskKey key, CompletableFuture<Void> future, Future<?> submission) {
+    final class MatterComponentTask {
+        private final MatterTaskKey key;
+        private final CompletableFuture<Void> future = new CompletableFuture<>();
+        private volatile Future<?> submission;
+        private volatile Throwable failure;
+        private boolean started;
+
+        MatterComponentTask(MatterTaskKey key) {
+            this.key = key;
+        }
+
+        synchronized boolean start() {
+            if (started || future.isDone()) {
+                return false;
+            }
+            started = true;
+            return true;
+        }
+
+        synchronized void setSubmission(Future<?> submission) {
+            this.submission = submission;
+            if (failure != null) {
+                submission.cancel(true);
+            }
+        }
+
+        synchronized boolean cancel(Throwable cause) {
+            if (future.isDone() || failure != null) {
+                return false;
+            }
+            failure = cause;
+            if (!started) {
+                finish(null);
+            }
+            if (submission != null) {
+                submission.cancel(true);
+            }
+            return true;
+        }
+
+        synchronized void finish(Throwable cause) {
+            if (failure == null) {
+                failure = cause;
+            } else if (cause != null && cause != failure) {
+                failure.addSuppressed(cause);
+            }
+            IN_FLIGHT_COMPONENTS.remove(key, this);
+            if (failure == null) {
+                future.complete(null);
+            } else {
+                future.completeExceptionally(failure);
+            }
+        }
+
+        void await(long timeoutMillis) {
+            long start = System.nanoTime();
+            boolean interrupted = false;
+            try {
+                while (true) {
+                    try {
+                        future.get(COMPONENT_TASK_POLL_MS, TimeUnit.MILLISECONDS);
+                        return;
+                    } catch (InterruptedException interruption) {
+                        // The writer cannot close while this task still writes through it, so the wait is
+                        // uninterruptible and the flag is restored on the way out.
+                        interrupted = true;
+                    } catch (TimeoutException timeout) {
+                        if (future.isDone() || failure != null) {
+                            continue;
+                        }
+                        Future<?> submitted = submission;
+                        boolean dropped = submitted != null && submitted.isDone();
+                        long waited = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start);
+                        if (!dropped && waited < timeoutMillis) {
+                            continue;
+                        }
+                        IllegalStateException timeoutFailure = new IllegalStateException("Mantle component " + key
+                                + (dropped ? " was dropped by the dispatcher" : " did not complete in " + waited + "ms"));
+                        if (cancel(timeoutFailure)) {
+                            IrisLogging.error(timeoutFailure.getMessage());
+                        }
+                    } catch (ExecutionException executionFailure) {
+                        Throwable cause = executionFailure.getCause();
+                        throw new CompletionException(cause != null ? cause : executionFailure);
+                    }
+                }
+            } finally {
+                if (interrupted) {
+                    Thread.currentThread().interrupt();
+                }
+            }
+        }
     }
 
     record MatterPassPlan(MantlePass pass, int passChunkRadius, int downstreamBlockRadius) {