Allow cancellation of image-heap laying out

ivan-ristovic · ivan-ristovic · commit 79efd4d30495 · 2025-05-20T12:01:35.000+02:00
diff --git a/substratevm/src/com.oracle.svm.core.genscavenge/src/com/oracle/svm/core/genscavenge/ChunkedImageHeapLayouter.java b/substratevm/src/com.oracle.svm.core.genscavenge/src/com/oracle/svm/core/genscavenge/ChunkedImageHeapLayouter.java
@@ -39,8 +39,10 @@
 import com.oracle.svm.core.genscavenge.remset.RememberedSet;
 import com.oracle.svm.core.hub.DynamicHub;
 import com.oracle.svm.core.image.ImageHeap;
-import com.oracle.svm.core.image.ImageHeapLayoutInfo;
 import com.oracle.svm.core.image.ImageHeapLayouter;
+import com.oracle.svm.core.image.ImageHeapLayingOutControl;
+import com.oracle.svm.core.image.ImageHeapLayingOutController;
+import com.oracle.svm.core.image.ImageHeapLayoutInfo;
 import com.oracle.svm.core.image.ImageHeapObject;
 import com.oracle.svm.core.option.SubstrateOptionsParser;
 import com.oracle.svm.core.util.UserError;
@@ -156,11 +158,12 @@ private Error reportHugeObjectError(ImageHeapObject info, String objectTypeMsg,
     }
 
     @Override
-    public ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize) {
+    public ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize, ImageHeapLayingOutController controller) {
+        ImageHeapLayingOutControl control = new ImageHeapLayingOutControl(controller);
         int objectAlignment = ConfigurationValues.getObjectLayout().getAlignment();
         assert pageSize % objectAlignment == 0 : "Page size does not match object alignment";
 
-        ImageHeapLayoutInfo layoutInfo = doLayout(imageHeap, pageSize);
+        ImageHeapLayoutInfo layoutInfo = doLayout(imageHeap, pageSize, control);
 
         for (ChunkedImageHeapPartition partition : getPartitions()) {
             assert partition.getStartOffset() % partition.getStartAlignment() == 0 : partition;
@@ -170,15 +173,16 @@ public ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize) {
         return layoutInfo;
     }
 
-    private ImageHeapLayoutInfo doLayout(ImageHeap imageHeap, int pageSize) {
+    private ImageHeapLayoutInfo doLayout(ImageHeap imageHeap, int pageSize, ImageHeapLayingOutControl control) {
         allocator = new ChunkedImageHeapAllocator(imageHeap, startOffset);
         for (ChunkedImageHeapPartition partition : getPartitions()) {
-            partition.layout(allocator);
+            control.heartbeat();
+            partition.layout(allocator, control);
         }
-        return populateInfoObjects(imageHeap.countAndVerifyDynamicHubs(), pageSize);
+        return populateInfoObjects(imageHeap.countAndVerifyDynamicHubs(), pageSize, control);
     }
 
-    private ImageHeapLayoutInfo populateInfoObjects(int dynamicHubCount, int pageSize) {
+    private ImageHeapLayoutInfo populateInfoObjects(int dynamicHubCount, int pageSize, ImageHeapLayingOutControl control) {
         // Determine writable start boundary from chunks: a chunk that contains writable objects
         // must also have a writable card table
         long offsetOfFirstWritableAlignedChunk = -1;
@@ -187,6 +191,7 @@ private ImageHeapLayoutInfo populateInfoObjects(int dynamicHubCount, int pageSiz
                 offsetOfFirstWritableAlignedChunk = chunk.getBegin();
                 break; // (chunks are in ascending memory order)
             }
+            control.heartbeat();
         }
         VMError.guarantee(offsetOfFirstWritableAlignedChunk >= 0 && offsetOfFirstWritableAlignedChunk % pageSize == 0, "Start of the writable part is assumed to be page-aligned");
         long offsetOfFirstWritableUnalignedChunk = -1;
@@ -199,6 +204,7 @@ private ImageHeapLayoutInfo populateInfoObjects(int dynamicHubCount, int pageSiz
                 offsetOfFirstWritableUnalignedChunk = chunk.getBegin();
             }
             offsetOfLastWritableUnalignedChunk = chunk.getBegin();
+            control.heartbeat();
         }
 
         heapInfo.initialize(getReadOnlyRegular().firstObject, getReadOnlyRegular().lastObject, getReadOnlyRelocatable().firstObject, getReadOnlyRelocatable().lastObject,
@@ -207,6 +213,8 @@ private ImageHeapLayoutInfo populateInfoObjects(int dynamicHubCount, int pageSiz
                         getReadOnlyHuge().firstObject, getReadOnlyHuge().lastObject, offsetOfFirstWritableAlignedChunk, offsetOfFirstWritableUnalignedChunk, offsetOfLastWritableUnalignedChunk,
                         dynamicHubCount);
 
+        control.heartbeat();
+
         long writableEnd = getWritableHuge().getStartOffset() + getWritableHuge().getSize();
         long writableSize = writableEnd - offsetOfFirstWritableAlignedChunk;
         long imageHeapSize = getReadOnlyHuge().getStartOffset() + getReadOnlyHuge().getSize() - startOffset;
diff --git a/substratevm/src/com.oracle.svm.core.genscavenge/src/com/oracle/svm/core/genscavenge/ChunkedImageHeapPartition.java b/substratevm/src/com.oracle.svm.core.genscavenge/src/com/oracle/svm/core/genscavenge/ChunkedImageHeapPartition.java
@@ -32,6 +32,7 @@
 import java.util.TreeMap;
 
 import com.oracle.svm.core.config.ConfigurationValues;
+import com.oracle.svm.core.image.ImageHeapLayingOutControl;
 import com.oracle.svm.core.image.ImageHeapObject;
 import com.oracle.svm.core.image.ImageHeapPartition;
 
@@ -72,47 +73,50 @@ void assign(ImageHeapObject obj) {
         objects.add(obj);
     }
 
-    void layout(ChunkedImageHeapAllocator allocator) {
+    void layout(ChunkedImageHeapAllocator allocator, ImageHeapLayingOutControl control) {
         if (hugeObjects) {
-            layoutInUnalignedChunks(allocator);
+            layoutInUnalignedChunks(allocator, control);
         } else {
-            layoutInAlignedChunks(allocator);
+            layoutInAlignedChunks(allocator, control);
         }
     }
 
-    private void layoutInUnalignedChunks(ChunkedImageHeapAllocator allocator) {
+    private void layoutInUnalignedChunks(ChunkedImageHeapAllocator allocator, ImageHeapLayingOutControl control) {
         allocator.finishAlignedChunk();
         allocator.alignBetweenChunks(getStartAlignment());
         startOffset = allocator.getPosition();
 
         for (ImageHeapObject info : objects) { // No need to sort by size
             appendAllocatedObject(info, allocator.allocateUnalignedChunkForObject(info, isWritable()));
+            control.heartbeat();
         }
 
         allocator.alignBetweenChunks(getEndAlignment());
         endOffset = allocator.getPosition();
     }
 
-    private void layoutInAlignedChunks(ChunkedImageHeapAllocator allocator) {
+    private void layoutInAlignedChunks(ChunkedImageHeapAllocator allocator, ImageHeapLayingOutControl control) {
         allocator.maybeStartAlignedChunk();
         allocator.alignInAlignedChunk(getStartAlignment());
         startOffset = allocator.getPosition();
 
-        allocateObjectsInAlignedChunks(allocator);
+        allocateObjectsInAlignedChunks(allocator, control);
 
         allocator.alignInAlignedChunk(getEndAlignment());
         endOffset = allocator.getPosition();
     }
 
-    private void allocateObjectsInAlignedChunks(ChunkedImageHeapAllocator allocator) {
+    private void allocateObjectsInAlignedChunks(ChunkedImageHeapAllocator allocator, ImageHeapLayingOutControl control) {
         NavigableMap<Long, Queue<ImageHeapObject>> sortedObjects = createSortedObjectsMap();
+        control.heartbeat();
         while (!sortedObjects.isEmpty()) {
             ImageHeapObject info = dequeueBestFit(sortedObjects, allocator.getRemainingBytesInAlignedChunk());
             if (info == null) {
                 allocator.startNewAlignedChunk();
             } else {
                 appendAllocatedObject(info, allocator.allocateObjectInAlignedChunk(info, isWritable()));
             }
+            control.heartbeat();
         }
     }
 
diff --git a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/ImageHeapLayingOutCancelledException.java b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/ImageHeapLayingOutCancelledException.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2025, 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.  Oracle designates this
+ * particular file as subject to the "Classpath" exception as provided
+ * by Oracle in the LICENSE file that accompanied this code.
+ *
+ * This code 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
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package com.oracle.svm.core.image;
+
+public class ImageHeapLayingOutCancelledException extends RuntimeException {
+    private static final long serialVersionUID = 1017980175582546348L;
+
+    public ImageHeapLayingOutCancelledException() {
+    }
+
+    public ImageHeapLayingOutCancelledException(RuntimeException cause) {
+        super(cause);
+    }
+}
diff --git a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/ImageHeapLayingOutControl.java b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/ImageHeapLayingOutControl.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) 2025, 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.  Oracle designates this
+ * particular file as subject to the "Classpath" exception as provided
+ * by Oracle in the LICENSE file that accompanied this code.
+ *
+ * This code 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
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package com.oracle.svm.core.image;
+
+/**
+ * Facilitates {@link ImageHeapLayouter#layout(ImageHeap, int, ImageHeapLayingOutController)}
+ * cancellation through an {@link ImageHeapLayingOutController} instance.
+ */
+public class ImageHeapLayingOutControl {
+    protected final ImageHeapLayingOutController controller;
+
+    public ImageHeapLayingOutControl(ImageHeapLayingOutController controller) {
+        this.controller = controller;
+    }
+
+    public void heartbeat() throws ImageHeapLayingOutCancelledException {
+        if (controller.shouldCancel()) {
+            throw new ImageHeapLayingOutCancelledException();
+        }
+    }
+}
diff --git a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/ImageHeapLayingOutController.java b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/ImageHeapLayingOutController.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2025, 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.  Oracle designates this
+ * particular file as subject to the "Classpath" exception as provided
+ * by Oracle in the LICENSE file that accompanied this code.
+ *
+ * This code 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
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package com.oracle.svm.core.image;
+
+public interface ImageHeapLayingOutController {
+
+    /**
+     * Called periodically by {@link ImageHeapLayingOutControl}.
+     *
+     * @return Whether the layout operation should be cancelled.
+     */
+    boolean shouldCancel();
+}
diff --git a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/ImageHeapLayouter.java b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/ImageHeapLayouter.java
@@ -45,7 +45,16 @@ public interface ImageHeapLayouter {
     /**
      * Places all heap partitions and assigns objects their final offsets.
      */
-    ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize);
+    default ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize) {
+        return layout(imageHeap, pageSize, new NonCancellableImageHeapLayingOutController());
+    }
+
+    /**
+     * Cancellable version of {@link ImageHeapLayouter#layout(ImageHeap, int)}. Cancellation is
+     * performed via {@link ImageHeapLayingOutController}.. If the layout is cancelled, an instance
+     * of {@link ImageHeapLayingOutCancelledException} is thrown.
+     */
+    ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize, ImageHeapLayingOutController controller);
 
     /** Hook to run tasks after heap layout is finished. */
     @SuppressWarnings("unused")
diff --git a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/NonCancellableImageHeapLayingOutController.java b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/NonCancellableImageHeapLayingOutController.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2025, 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.  Oracle designates this
+ * particular file as subject to the "Classpath" exception as provided
+ * by Oracle in the LICENSE file that accompanied this code.
+ *
+ * This code 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
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package com.oracle.svm.core.image;
+
+/**
+ * An implementation of {@link ImageHeapLayingOutController} that never cancels the laying out.
+ */
+public final class NonCancellableImageHeapLayingOutController implements ImageHeapLayingOutController {
+    @Override
+    public boolean shouldCancel() {
+        return false;
+    }
+}
diff --git a/web-image/src/com.oracle.svm.hosted.webimage/src/com/oracle/svm/hosted/webimage/wasmgc/image/WasmGCHeapLayouter.java b/web-image/src/com.oracle.svm.hosted.webimage/src/com/oracle/svm/hosted/webimage/wasmgc/image/WasmGCHeapLayouter.java
@@ -30,6 +30,8 @@
 
 import com.oracle.graal.pointsto.heap.ImageHeapPrimitiveArray;
 import com.oracle.svm.core.image.ImageHeap;
+import com.oracle.svm.core.image.ImageHeapLayingOutControl;
+import com.oracle.svm.core.image.ImageHeapLayingOutController;
 import com.oracle.svm.core.image.ImageHeapLayouter;
 import com.oracle.svm.core.image.ImageHeapObject;
 import com.oracle.svm.core.image.ImageHeapPartition;
@@ -73,33 +75,37 @@ public void assignObjectToPartition(ImageHeapObject info, boolean immutable, boo
     }
 
     @Override
-    public WasmGCImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize) {
-        layoutPseudoPartition();
-        doLayout();
+    public WasmGCImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize, ImageHeapLayingOutController controller) {
+        ImageHeapLayingOutControl control = new ImageHeapLayingOutControl(controller);
+
+        layoutPseudoPartition(control);
+        doLayout(control);
 
         long totalSize = StreamSupport.stream(imageHeap.getObjects().spliterator(), false).mapToLong(ImageHeapObject::getSize).sum();
         long serializedSize = singlePartition.getStartOffset() + singlePartition.getSize() - startOffset;
         return new WasmGCImageHeapLayoutInfo(startOffset, serializedSize, totalSize);
     }
 
-    private void doLayout() {
+    private void doLayout(ImageHeapLayingOutControl control) {
         int offset = 0;
         for (ImageHeapObject info : singlePartition.getObjects()) {
             // Only primitive arrays are supposed to be in this partition
             ImageHeapPrimitiveArray primitiveArray = (ImageHeapPrimitiveArray) info.getWrapped();
             info.setOffsetInPartition(offset);
             offset += primitiveArray.getType().getComponentType().getStorageKind().getByteCount() * primitiveArray.getLength();
+            control.heartbeat();
         }
 
         singlePartition.setSize(offset);
     }
 
-    private void layoutPseudoPartition() {
+    private void layoutPseudoPartition(ImageHeapLayingOutControl control) {
         // Approximate size of the partition (based on the SVM object layout size of the objects)
         long pseudoSize = 0;
         for (ImageHeapObject info : pseudoPartition.getObjects()) {
             info.setOffsetInPartition(0);
             pseudoSize = info.getSize();
+            control.heartbeat();
         }
 
         pseudoPartition.setSize(pseudoSize);
