Allow cancellation of image-heap laying out

Author: ivan-ristovic

substratevm/src/com.oracle.svm.core.genscavenge/src/com/oracle/svm/core/genscavenge/ChunkedImageHeapLayouter.java

@@ -39,8 +39,8 @@
 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.ImageHeapLayoutInfo;
 import com.oracle.svm.core.image.ImageHeapObject;
 import com.oracle.svm.core.option.SubstrateOptionsParser;
 import com.oracle.svm.core.util.UserError;
@@ -156,11 +156,12 @@ private Error reportHugeObjectError(ImageHeapObject info, String objectTypeMsg,
     }
 
     @Override
-    public ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize) {
+    public ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize, ImageHeapLayouterCallback callback) {
+        ImageHeapLayouterControl control = new ImageHeapLayouterControl(callback);
         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 +171,16 @@ public ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize) {
         return layoutInfo;
     }
 
-    private ImageHeapLayoutInfo doLayout(ImageHeap imageHeap, int pageSize) {
+    private ImageHeapLayoutInfo doLayout(ImageHeap imageHeap, int pageSize, ImageHeapLayouterControl control) {
         allocator = new ChunkedImageHeapAllocator(imageHeap, startOffset);
         for (ChunkedImageHeapPartition partition : getPartitions()) {
-            partition.layout(allocator);
+            control.poll();
+            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, ImageHeapLayouterControl control) {
         // Determine writable start boundary from chunks: a chunk that contains writable objects
         // must also have a writable card table
         long offsetOfFirstWritableAlignedChunk = -1;
@@ -188,6 +190,8 @@ private ImageHeapLayoutInfo populateInfoObjects(int dynamicHubCount, int pageSiz
                 break; // (chunks are in ascending memory order)
             }
         }
+        control.poll();
+
         VMError.guarantee(offsetOfFirstWritableAlignedChunk >= 0 && offsetOfFirstWritableAlignedChunk % pageSize == 0, "Start of the writable part is assumed to be page-aligned");
         long offsetOfFirstWritableUnalignedChunk = -1;
         long offsetOfLastWritableUnalignedChunk = -1;
@@ -200,13 +204,16 @@ private ImageHeapLayoutInfo populateInfoObjects(int dynamicHubCount, int pageSiz
             }
             offsetOfLastWritableUnalignedChunk = chunk.getBegin();
         }
+        control.poll();
 
         heapInfo.initialize(getReadOnlyRegular().firstObject, getReadOnlyRegular().lastObject, getReadOnlyRelocatable().firstObject, getReadOnlyRelocatable().lastObject,
                         getWritablePatched().firstObject, getWritablePatched().lastObject,
                         getWritableRegular().firstObject, getWritableRegular().lastObject, getWritableHuge().firstObject, getWritableHuge().lastObject,
                         getReadOnlyHuge().firstObject, getReadOnlyHuge().lastObject, offsetOfFirstWritableAlignedChunk, offsetOfFirstWritableUnalignedChunk, offsetOfLastWritableUnalignedChunk,
                         dynamicHubCount);
 
+        control.poll();
+
         long writableEnd = getWritableHuge().getStartOffset() + getWritableHuge().getSize();
         long writableSize = writableEnd - offsetOfFirstWritableAlignedChunk;
         long imageHeapSize = getReadOnlyHuge().getStartOffset() + getReadOnlyHuge().getSize() - startOffset;

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.ImageHeapLayouter.ImageHeapLayouterControl;
 import com.oracle.svm.core.image.ImageHeapObject;
 import com.oracle.svm.core.image.ImageHeapPartition;
 
@@ -72,44 +73,46 @@ void assign(ImageHeapObject obj) {
         objects.add(obj);
     }
 
-    void layout(ChunkedImageHeapAllocator allocator) {
+    void layout(ChunkedImageHeapAllocator allocator, ImageHeapLayouterControl control) {
         if (hugeObjects) {
-            layoutInUnalignedChunks(allocator);
+            layoutInUnalignedChunks(allocator, control);
         } else {
-            layoutInAlignedChunks(allocator);
+            layoutInAlignedChunks(allocator, control);
         }
     }
 
-    private void layoutInUnalignedChunks(ChunkedImageHeapAllocator allocator) {
+    private void layoutInUnalignedChunks(ChunkedImageHeapAllocator allocator, ImageHeapLayouterControl 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.poll();
         }
 
         allocator.alignBetweenChunks(getEndAlignment());
         endOffset = allocator.getPosition();
     }
 
-    private void layoutInAlignedChunks(ChunkedImageHeapAllocator allocator) {
+    private void layoutInAlignedChunks(ChunkedImageHeapAllocator allocator, ImageHeapLayouterControl 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, ImageHeapLayouterControl control) {
         NavigableMap<Long, Queue<ImageHeapObject>> sortedObjects = createSortedObjectsMap();
         while (!sortedObjects.isEmpty()) {
             ImageHeapObject info = dequeueBestFit(sortedObjects, allocator.getRemainingBytesInAlignedChunk());
             if (info == null) {
                 allocator.startNewAlignedChunk();
+                control.poll();
             } else {
                 appendAllocatedObject(info, allocator.allocateObjectInAlignedChunk(info, isWritable()));
             }

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/image/ImageHeapLayouter.java

@@ -43,9 +43,11 @@ public interface ImageHeapLayouter {
     void assignObjectToPartition(ImageHeapObject info, boolean immutable, boolean references, boolean relocatable, boolean patched);
 
     /**
-     * Places all heap partitions and assigns objects their final offsets.
+     * Places all heap partitions and assigns objects their final offsets. The layout operation can
+     * be cancelled through the {@link ImageHeapLayouterCallback}. If the layout is cancelled, an
+     * instance of {@link ImageHeapLayoutCancelledException} is thrown.
      */
-    ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize);
+    ImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize, ImageHeapLayouterCallback callback);
 
     /** Hook to run tasks after heap layout is finished. */
     @SuppressWarnings("unused")
@@ -63,4 +65,39 @@ default void afterLayout(ImageHeap imageHeap) {
      *            same offset in the given buffer, the same offset must be specified to this method.
      */
     void writeMetadata(ByteBuffer imageHeapBytes, long imageHeapOffsetInBuffer);
+
+    /**
+     * Facilitates {@link ImageHeapLayouter#layout} cancellation through an
+     * {@link ImageHeapLayouterCallback} instance.
+     */
+    class ImageHeapLayouterControl {
+        protected final ImageHeapLayouterCallback callback;
+
+        public ImageHeapLayouterControl(ImageHeapLayouterCallback callback) {
+            this.callback = callback;
+        }
+
+        public void poll() throws ImageHeapLayoutCancelledException {
+            if (callback.shouldCancel()) {
+                throw new ImageHeapLayoutCancelledException();
+            }
+        }
+    }
+
+    interface ImageHeapLayouterCallback {
+
+        ImageHeapLayouterCallback NONE = () -> false;
+
+        /**
+         * Called periodically to determine whether the operation should be canceled.
+         */
+        boolean shouldCancel();
+    }
+
+    class ImageHeapLayoutCancelledException extends RuntimeException {
+        private static final long serialVersionUID = 1017980175582546348L;
+
+        public ImageHeapLayoutCancelledException() {
+        }
+    }
 }

substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/image/NativeImage.java

@@ -94,6 +94,7 @@
 import com.oracle.svm.core.graal.nodes.TLABObjectHeaderConstant;
 import com.oracle.svm.core.heap.Heap;
 import com.oracle.svm.core.image.ImageHeapLayoutInfo;
+import com.oracle.svm.core.image.ImageHeapLayouter.ImageHeapLayouterCallback;
 import com.oracle.svm.core.image.ImageHeapPartition;
 import com.oracle.svm.core.imagelayer.DynamicImageLayerInfo;
 import com.oracle.svm.core.imagelayer.ImageLayerBuildingSupport;
@@ -436,7 +437,7 @@ public void build(String imageName, DebugContext debug) {
 
             long roSectionSize = codeCache.getAlignedConstantsSize();
             long rwSectionSize = ConfigurationValues.getObjectLayout().alignUp(cGlobals.getSize());
-            heapLayout = heap.getLayouter().layout(heap, objectFile.getPageSize());
+            heapLayout = heap.getLayouter().layout(heap, objectFile.getPageSize(), ImageHeapLayouterCallback.NONE);
             // after this point, the layout is final and must not be changed anymore
             assert !hasDuplicatedObjects(heap.getObjects()) : "heap.getObjects() must not contain any duplicates";
 

web-image/src/com.oracle.svm.hosted.webimage/src/com/oracle/svm/hosted/webimage/wasm/codegen/WebImageWasmLMCodeGen.java

@@ -36,6 +36,7 @@
 import com.oracle.svm.core.c.CGlobalData;
 import com.oracle.svm.core.graal.code.CGlobalDataReference;
 import com.oracle.svm.core.image.ImageHeapLayoutInfo;
+import com.oracle.svm.core.image.ImageHeapLayouter.ImageHeapLayouterCallback;
 import com.oracle.svm.core.meta.MethodPointer;
 import com.oracle.svm.hosted.image.NativeImageHeap;
 import com.oracle.svm.hosted.image.NativeImageHeapWriter;
@@ -90,7 +91,7 @@ public long getImageHeapSize() {
      */
     @Override
     protected void writeImageHeap() {
-        ImageHeapLayoutInfo layout = codeCache.nativeImageHeap.getLayouter().layout(codeCache.nativeImageHeap, WasmUtil.PAGE_SIZE);
+        ImageHeapLayoutInfo layout = codeCache.nativeImageHeap.getLayouter().layout(codeCache.nativeImageHeap, WasmUtil.PAGE_SIZE, ImageHeapLayouterCallback.NONE);
         setLayout(layout);
 
         afterHeapLayout();

web-image/src/com.oracle.svm.hosted.webimage/src/com/oracle/svm/hosted/webimage/wasmgc/codegen/WasmGCHeapWriter.java

@@ -49,6 +49,7 @@
 import com.oracle.graal.pointsto.util.AnalysisError;
 import com.oracle.svm.core.hub.DynamicHub;
 import com.oracle.svm.core.hub.Hybrid;
+import com.oracle.svm.core.image.ImageHeapLayouter.ImageHeapLayouterCallback;
 import com.oracle.svm.core.meta.MethodPointer;
 import com.oracle.svm.core.meta.SubstrateMethodPointerConstant;
 import com.oracle.svm.core.util.VMError;
@@ -273,7 +274,7 @@ public WasmGCHeapWriter(WebImageCodeCache codeCache, WebImageWasmGCProviders pro
 
     public WasmGCImageHeapLayoutInfo layout() {
         collectObjectData();
-        return (WasmGCImageHeapLayoutInfo) heap.getLayouter().layout(heap, WasmUtil.PAGE_SIZE);
+        return (WasmGCImageHeapLayoutInfo) heap.getLayouter().layout(heap, WasmUtil.PAGE_SIZE, ImageHeapLayouterCallback.NONE);
     }
 
     public void write(WasmGCImageHeapLayoutInfo layout, WasmModule module) {

web-image/src/com.oracle.svm.hosted.webimage/src/com/oracle/svm/hosted/webimage/wasmgc/image/WasmGCHeapLayouter.java

@@ -73,7 +73,7 @@ public void assignObjectToPartition(ImageHeapObject info, boolean immutable, boo
     }
 
     @Override
-    public WasmGCImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize) {
+    public WasmGCImageHeapLayoutInfo layout(ImageHeap imageHeap, int pageSize, ImageHeapLayouterCallback callback) {
         layoutPseudoPartition();
         doLayout();