[GR-64947] Implement Engine.storeCache(...) method to persist engine caches manually.

sdk/CHANGELOG.md

@@ -20,6 +20,7 @@ This changelog summarizes major changes between GraalVM SDK versions. The main f
 * GR-66339 Truffle now checks that its multi-release classes ([JEP 238](https://openjdk.org/jeps/238)) are loaded correctly and provides a descriptive error message if they are not.
 * GR-55996 Added the options `engine.SourceCacheStatistics` and `engine.SourceCacheStatisticDetails` to print polyglot source cache statistics on engine close.
 
+* GR-64947 Added `Engine.storeCache(Path)` to manually store [auxiliary engine caches](https://github.com/oracle/graal/blob/master/truffle/docs/AuxiliaryEngineCachingEnterprise.md) when needed.
 
 ## Version 24.2.0
 * GR-54905 When using Truffle NFI with the Panama backend, native access must now be granted to the Truffle module instead of the NFI Panama module. Use the `--enable-native-access=org.graalvm.truffle` Java command line option to enable the native access for the NFI Panama backend.

sdk/src/org.graalvm.polyglot/snapshot.sigtest

@@ -210,6 +210,8 @@ intf java.lang.AutoCloseable
 meth public !varargs static boolean copyResources(java.nio.file.Path,java.lang.String[]) throws java.io.IOException
 meth public !varargs static org.graalvm.polyglot.Engine create(java.lang.String[])
 meth public !varargs static org.graalvm.polyglot.Engine$Builder newBuilder(java.lang.String[])
+meth public boolean storeCache(java.nio.file.Path)
+meth public boolean storeCache(java.nio.file.Path,org.graalvm.nativeimage.c.type.WordPointer)
 meth public java.lang.String getImplementationName()
 meth public java.lang.String getVersion()
 meth public java.util.Map<java.lang.String,org.graalvm.polyglot.Instrument> getInstruments()

sdk/src/org.graalvm.polyglot/src/org/graalvm/polyglot/Engine.java

@@ -80,6 +80,7 @@
 import java.util.Properties;
 import java.util.ServiceLoader;
 import java.util.Set;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -90,6 +91,7 @@
 import org.graalvm.home.HomeFinder;
 import org.graalvm.home.Version;
 import org.graalvm.nativeimage.ImageInfo;
+import org.graalvm.nativeimage.c.type.WordPointer;
 import org.graalvm.options.OptionDescriptor;
 import org.graalvm.options.OptionDescriptors;
 import org.graalvm.polyglot.HostAccess.MutableTargetMapping;
@@ -316,6 +318,76 @@ public void close() {
         close(false);
     }
 
+    /**
+     * Stores the auxiliary engine cache to the targetFile without cancellation.
+     *
+     * @see #storeCache(Path, WordPointer)
+     * @throws UnsupportedOperationException if this engine or the host virtual machine does not
+     *             support storing the cache.
+     * @since 25.0
+     */
+    public boolean storeCache(Path targetFile) throws UnsupportedOperationException {
+        return dispatch.storeCache(receiver, targetFile, 0L);
+    }
+
+    /**
+     * Stores the auxiliary engine cache to the {@code targetFile}. If it already exists, the file
+     * will be overwritten. The option <code>engine.CacheStoreEnabled</code> must be set to
+     * <code>true</code> to use this feature. Stored caches may be loaded by specifying the path
+     * using the <code>engine.CacheLoad</code> option.
+     * <p>
+     * Note that this feature is experimental and only supported on native-image hosts with
+     * Truffle's enterprise extensions.
+     * </p>
+     *
+     * <h3>Basic Usage:</h3>
+     *
+     * <pre>
+     * // Store the engine cache into a file
+     * Path store = Files.createTempFile("cache", "engine");
+     * try (Engine e = Engine.newBuilder().allowExperimentalOptions(true).option("engine.CacheStoreEnabled", "true").build()) {
+     *     try (Context c = Context.newBuilder().engine(e).build()) {
+     *         // Evaluate sources, run application
+     *     }
+     *     e.storeCache(store);
+     * }
+     *
+     * // Load the engine cache from a file
+     * try (Engine e = Engine.newBuilder().allowExperimentalOptions(true).option("engine.CacheLoad", store.toAbsolutePath().toString()).build()) {
+     *     try (Context c = Context.newBuilder().engine(e).build()) {
+     *         // The context should be able to use
+     *         // the existing code cache.
+     *     }
+     * }
+     * </pre>
+     *
+     * <p>
+     * See the <a href=
+     * "https://github.com/oracle/graal/blob/master/truffle/docs/AuxiliaryEngineCachingEnterprise.md">
+     * documentation</a> on auxiliary engine caching for further details.
+     * </p>
+     *
+     * @param targetFile the file to which the cache is stored
+     * @param cancelledWord a native pointer; if set to a non-zero value, the operation is
+     *            cancelled. Allows cancellation of the cache store operation through a
+     *            <code>cancelled</code> control word. The memory {@code address} pointing to the
+     *            control word is polled periodically during storage without guaranteed frequency
+     *            and may be delayed by safepoints such as garbage collection. A control word value
+     *            of zero must be maintained for the duration of the operation. If a non-zero value
+     *            is detected, the operation will be cancelled. A non-null provided pointer must
+     *            remain accessible during the entire operation. Providing an invalid or
+     *            inaccessible pointer may result in a VM crash.
+     * @return <code>true</code> if the file was written; otherwise, <code>false</code>
+     * @throws CancellationException if the storeCache operation was cancelled via the
+     *             <code>cancelled</code> pointer
+     * @throws UnsupportedOperationException if this engine or host virtual machine does not support
+     *             cache storage
+     * @since 25.0
+     */
+    public boolean storeCache(Path targetFile, WordPointer cancelledWord) throws CancellationException, UnsupportedOperationException {
+        return dispatch.storeCache(receiver, targetFile, cancelledWord.rawValue());
+    }
+
     /**
      * Gets a human-readable name of the polyglot implementation (for example, "Default Truffle
      * Engine" or "Graal Truffle Engine"). The returned value may change without notice. The value

sdk/src/org.graalvm.polyglot/src/org/graalvm/polyglot/impl/AbstractPolyglotImpl.java

@@ -852,6 +852,8 @@ public abstract Object attachExecutionListener(Object engine, Consumer<Object> o
 
         public abstract void onEngineCollected(Object engineReceiver);
 
+        public abstract boolean storeCache(Object engineReceiver, Path targetFile, long cancelledWord);
+
     }
 
     public abstract static class AbstractExceptionDispatch extends AbstractDispatchClass {

truffle/src/com.oracle.truffle.api.test/src/com/oracle/truffle/api/test/wrapper/HostEngineDispatch.java

@@ -43,6 +43,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.lang.ref.Reference;
+import java.nio.file.Path;
 import java.time.ZoneId;
 import java.util.Map;
 import java.util.Set;
@@ -188,4 +189,10 @@ public void onEngineCollected(Object receiver) {
         dispatch.onEngineCollected(engineReceiver);
         hostToGuest.shutdown(engine.remoteEngine);
     }
+
+    @Override
+    public boolean storeCache(Object engineReceiver, Path targetFile, long cancelledWord) {
+        throw new UnsupportedOperationException();
+    }
+
 }

truffle/src/com.oracle.truffle.api/src/com/oracle/truffle/api/impl/Accessor.java

@@ -694,7 +694,9 @@ public abstract <T, G> Iterator<T> mergeHostGuestFrames(Object polyglotEngine, S
 
         public abstract void preinitializeContext(Object polyglotEngine);
 
-        public abstract void finalizeStore(Object polyglotEngine);
+        public abstract Object finalizeStore(Object polyglotEngine);
+
+        public abstract void restoreStore(Object polyglotEngine, Object finalizationResult);
 
         public abstract Object getEngineLock(Object polyglotEngine);
 
@@ -1297,6 +1299,8 @@ public final FrameExtensions getFrameExtensionsUnsafe() {
 
         public abstract boolean onEngineClosing(Object runtimeData);
 
+        public abstract boolean onStoreCache(Object runtimeData, Path targetPath, long cancelledWord);
+
         public abstract void onEngineClosed(Object runtimeData);
 
         public abstract boolean isOSRRootNode(RootNode rootNode);

truffle/src/com.oracle.truffle.api/src/com/oracle/truffle/api/impl/DefaultRuntimeAccessor.java

@@ -40,6 +40,7 @@
  */
 package com.oracle.truffle.api.impl;
 
+import java.nio.file.Path;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -255,6 +256,11 @@ public boolean onEngineClosing(Object runtimeData) {
             return false;
         }
 
+        @Override
+        public boolean onStoreCache(Object runtimeData, Path targetPath, long cancelledWord) {
+            throw new UnsupportedOperationException("Persisting an engine is not supported with the the Truffle fallback runtime. It is only supported on native-image hosts.");
+        }
+
         @Override
         public boolean isOSRRootNode(RootNode rootNode) {
             return false;

truffle/src/com.oracle.truffle.polyglot/src/com/oracle/truffle/polyglot/EngineAccessor.java

@@ -130,6 +130,7 @@
 import com.oracle.truffle.api.source.SourceSection;
 import com.oracle.truffle.polyglot.FileSystems.ResetablePath;
 import com.oracle.truffle.polyglot.PolyglotContextConfig.FileSystemConfig;
+import com.oracle.truffle.polyglot.PolyglotEngineImpl.FinalizationResult;
 import com.oracle.truffle.polyglot.PolyglotImpl.EmbedderFileSystemContext;
 import com.oracle.truffle.polyglot.PolyglotImpl.VMObject;
 import com.oracle.truffle.polyglot.PolyglotLocals.InstrumentContextLocal;
@@ -1323,8 +1324,13 @@ public void preinitializeContext(Object polyglotEngine) {
         }
 
         @Override
-        public void finalizeStore(Object polyglotEngine) {
-            ((PolyglotEngineImpl) polyglotEngine).finalizeStore();
+        public Object finalizeStore(Object polyglotEngine) {
+            return ((PolyglotEngineImpl) polyglotEngine).finalizeStore();
+        }
+
+        @Override
+        public void restoreStore(Object polyglotEngine, Object finalizationResult) {
+            ((PolyglotEngineImpl) polyglotEngine).restoreStore((FinalizationResult) finalizationResult);
         }
 
         @Override

truffle/src/com.oracle.truffle.polyglot/src/com/oracle/truffle/polyglot/PolyglotEngineDispatch.java

@@ -43,6 +43,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.lang.ref.Reference;
+import java.nio.file.Path;
 import java.time.ZoneId;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -327,4 +328,11 @@ public SandboxPolicy getSandboxPolicy(Object engineReceiver) {
     public void onEngineCollected(Object engineReceiver) {
         ((PolyglotEngineImpl) engineReceiver).onEngineCollected();
     }
+
+    @Override
+    public boolean storeCache(Object engineReceiver, Path targetFile, long cancelledWord) {
+        PolyglotEngineImpl engine = ((PolyglotEngineImpl) engineReceiver);
+        return engine.storeCache(targetFile, cancelledWord);
+    }
+
 }

truffle/src/com.oracle.truffle.polyglot/src/com/oracle/truffle/polyglot/PolyglotEngineImpl.java

@@ -1249,6 +1249,28 @@ <T extends TruffleLanguage<?>> PolyglotLanguage getLanguage(Class<T> languageCla
         return foundLanguage;
     }
 
+    boolean storeCache(Path targetPath, long cancelledWord) {
+        if (!TruffleOptions.AOT) {
+            throw new UnsupportedOperationException("Storing the engine cache is only supported on native-image hosts.");
+        }
+
+        synchronized (this.lock) {
+            if (closingThread != null || closed) {
+                throw new IllegalStateException("The engine is already closed and cannot be cancelled or persisted.");
+            }
+            if (!storeEngine) {
+                throw new IllegalStateException(
+                                "In order to store the cache the option 'engine.CacheStoreEnabled' must be set to 'true'.");
+            }
+            List<PolyglotContextImpl> localContexts = collectAliveContexts();
+            if (!localContexts.isEmpty()) {
+                throw new IllegalStateException("There are still alive contexts that need to be closed or cancelled before the engine can be persisted.");
+            }
+
+            return RUNTIME.onStoreCache(this.runtimeData, targetPath, cancelledWord);
+        }
+    }
+
     void ensureClosed(boolean force, boolean initiatedByContext) {
         synchronized (this.lock) {
             Thread currentThread = Thread.currentThread();
@@ -1525,12 +1547,15 @@ void preInitialize() {
         }
     }
 
+    record FinalizationResult(DispatchOutputStream out, DispatchOutputStream err, InputStream in) {
+    }
+
     /**
      * Invoked when the context is closing to prepare an engine to be stored.
      */
-    void finalizeStore() {
+    FinalizationResult finalizeStore() {
         assert Thread.holdsLock(this.lock);
-
+        FinalizationResult result = new FinalizationResult(out, err, in);
         this.out = null;
         this.err = null;
         this.in = null;
@@ -1544,6 +1569,13 @@ void finalizeStore() {
         if (hostLanguageService != null) {
             hostLanguageService.release();
         }
+        return result;
+    }
+
+    void restoreStore(FinalizationResult result) {
+        this.out = result.out;
+        this.err = result.err;
+        this.in = result.in;
     }
 
     @TruffleBoundary

truffle/src/com.oracle.truffle.polyglot/src/com/oracle/truffle/polyglot/PolyglotImpl.java

@@ -116,6 +116,7 @@ public final class PolyglotImpl extends AbstractPolyglotImpl {
                     "engine.Cache",
                     "engine.CacheLoad",
                     "engine.CacheStore",
+                    "engine.CacheStoreEnabled",
                     "engine.DebugCacheLoad",
                     "engine.DebugCacheStore",
                     "engine.SpawnIsolate");

truffle/src/com.oracle.truffle.runtime/src/com/oracle/truffle/runtime/CompilationTask.java

@@ -50,6 +50,7 @@
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
 
 import com.oracle.truffle.api.Truffle;
@@ -70,13 +71,15 @@ public void accept(CompilationTask task) {
             }
         }
     };
+
     final WeakReference<OptimizedCallTarget> targetRef;
     private final BackgroundCompileQueue.Priority priority;
     private final long id;
     private final Consumer<CompilationTask> action;
     private final EngineData engineData;
     private volatile Future<?> future;
     private volatile boolean cancelled;
+    volatile BooleanSupplier cancelledPredicate;
     private volatile boolean started;
     // Traversing queue related
     private int lastCount;
@@ -105,10 +108,12 @@ private CompilationTask(BackgroundCompileQueue.Priority priority, WeakReference<
             lastCount = Integer.MIN_VALUE;
             engineData = null;
             isOSR = false;
+            cancelledPredicate = null;
         } else {
             lastCount = target.getCallAndLoopCount();
             engineData = target.engine;
             isOSR = target.isOSR();
+            cancelledPredicate = target.engine.cancelledPredicate;
         }
     }
 
@@ -170,7 +175,8 @@ public synchronized boolean start() {
 
     @Override
     public boolean isCancelled() {
-        return cancelled;
+        BooleanSupplier cancelPredicate = cancelledPredicate;
+        return cancelled || (cancelPredicate != null && cancelPredicate.getAsBoolean());
     }
 
     @Override

truffle/src/com.oracle.truffle.runtime/src/com/oracle/truffle/runtime/EngineCacheSupport.java

@@ -40,6 +40,7 @@
  */
 package com.oracle.truffle.runtime;
 
+import java.nio.file.Path;
 import java.util.function.Function;
 
 import org.graalvm.options.OptionDescriptors;
@@ -55,6 +56,11 @@ public interface EngineCacheSupport extends OptimizedRuntimeServiceProvider {
 
     boolean onEngineClosing(EngineData e);
 
+    @SuppressWarnings("unused")
+    default boolean onStoreCache(EngineData e, Path path, long cancelledWord) {
+        throw new UnsupportedOperationException("Engine persist ist not yet supported on this JDK. Please update to resolve this problem.");
+    }
+
     void onEngineClosed(EngineData e);
 
     boolean isStoreEnabled(OptionValues options);