Save the JIT's generated code in warm snapshots

docs/browser.md

@@ -160,9 +160,11 @@ With `warmSnapshot: true`, `cheerpjRunLibrary` also saves its VM after a call
 returns once its calls have loaded at least 200 classes since boot (one javac
 compile loads over a thousand) and no library handles are live. Later library
 starts with the same class path restore that image, after checking that the
-files, missing paths and directory listings those calls used are unchanged; the
-first javac compile after it took 2.7 s instead of 5.3 s in Chromium on a
-Raspberry Pi 5. An image is about 18 MB. IndexedDB keeps the 2 most recent warm images separately
+files, missing paths and directory listings those calls used are unchanged. In
+Chromium, the first javac compile after it takes about half the time of one
+after a boot snapshot. An image is about 21 MB; it holds the JS the JIT
+generated for the methods compiled so far, and a restored method compiles that
+code on its first call. IndexedDB keeps the 2 most recent warm images separately
 from boot snapshots. With `debug: true`, the console says whether a VM was
 saved, skipped as busy (with the reason) or restored, and why a restore missed.
 

docs/node.md

@@ -74,8 +74,10 @@ returns once its calls have loaded at least 200 classes since boot (one javac
 compile loads over a thousand) and no library handles are live. A later library
 start with the same class path and working directory restores it, after
 checking that the files, missing paths and directory listings those calls used
-are unchanged; a javac compile from it takes about 1 s instead of 6 s on a
-Raspberry Pi 5. An image is about 18 MB; warm snapshots share the snapshot
+are unchanged. A first javac compile from it takes about a sixth of the time
+of one after a boot snapshot. An image is about 21 MB, as it also holds the JS
+the JIT generated for the methods compiled so far; a restored method compiles
+that code on its first call instead of generating it again. Warm snapshots share the snapshot
 directory and keep the 2 most recent. They build on boot snapshots, so with
 `snapshotDir: null` there are none.
 

src/jit.ts

@@ -22,6 +22,7 @@ import { memberTarget } from "./natives/invoke.ts";
 import { concatStrings } from "./natives/strings.ts";
 import { Frame, JavaThrow, NativeFrame, PendingThrow, fromNative } from "./thread.ts";
 import type { JThread } from "./thread.ts";
+import type { JVM } from "./jvm.ts";
 
 export const RET = 0, SUSPEND = 1;
 export type Compiled = (t: JThread, f: Frame) => number;
@@ -64,10 +65,26 @@ export function noteBackedge(t: JThread, m: JMethod) {
   if (m.hot >= BACKEDGE_THRESHOLD && m.jit === null && !m.noJit) compile(t, m);
 }
 
+export interface SavedJit { src: string; consts: any[] }
+
+// Gives m the code saved with a warm snapshot, compiled on its first call.
+export function restoreJit(vm: JVM, m: JMethod, code: SavedJit) {
+  if (disabled || noInline || stats) return;
+  vm.jitCode.set(m, code);
+  m.hot = THRESHOLD - 1;
+}
+
 function compile(t: JThread, m: JMethod) {
   if (disabled || m.noJit || !m.code || m.name === "<clinit>" || m.code.code.length > MAX_CODE) return;
+  const vm = t.vm;
   try {
-    const { src, consts } = generate(t, m);
+    let code = vm.keepJit ? undefined : vm.jitCode.get(m);
+    if (code) vm.jitCode.delete(m);
+    else {
+      code = generate(t, m);
+      if (vm.keepJit && !stats && !noInline) vm.jitCode.set(m, code);
+    }
+    const { src, consts } = code;
     if (dump && m.toString().includes(dump)) console.error(`// ${m}\n${src}`);
     m.jit = new Function("rt", "K", src)(runtime, consts) as Compiled;
     if (stats) stats.set(`compiled ${m}`, src.length);

src/jvm.ts

@@ -25,6 +25,7 @@ import { hostNative, jniNames } from "./hostnatives.ts";
 import { Library } from "./library.ts";
 import { dumpVM, restoreVM, snapshotKey } from "./snapshot.ts";
 import type { SnapshotImage, WarmSnapshotImage } from "./snapshot.ts";
+import type { SavedJit } from "./jit.ts";
 import type { LibraryRequest } from "./library.ts";
 import type { HostNative } from "./hostnatives.ts";
 import { PROXY_HOST, PROXY_PORT } from "./http-proxy.ts";
@@ -143,6 +144,10 @@ export class JVM {
   classCount = 0;
   // every class in creation order, until the snapshot after boot is taken
   created: JClass[] | null;
+  // Generated JS of compiled methods: kept until the warm snapshot is saved,
+  // then (after a warm restore) compiled on each method's first call.
+  keepJit: boolean;
+  jitCode = new Map<JMethod, SavedJit>();
   threadsStarted = 0;
   peakThreads = 0;
   exitCode: number | null = null;
@@ -159,6 +164,7 @@ export class JVM {
 
   constructor(opts: JvmOptions) {
     this.created = opts.snapshots ? [] : null;
+    this.keepJit = !!(opts.warmSnapshot && opts.warmSnapshots && opts.snapshots);
     this.opts = opts;
     this.fs = opts.fs;
     const extra = ipv4Literal(opts.hostName);
@@ -1139,10 +1145,17 @@ export class JVM {
       this.warmTraceActive = this.fs.beginWarmTrace();
       this.warmTraceClasses = this.classCount;
     }
+    // no warm snapshot will be taken (a restored VM keeps its saved code)
+    if (!this.warmTraceActive && this.keepJit) this.dropJit();
     started();
     yield* this.library.loop(t);
   }
 
+  private dropJit() {
+    this.keepJit = false;
+    this.jitCode.clear();
+  }
+
   libraryIdle(t: JThread) {
     const cache = this.opts.warmSnapshots;
     if (!this.opts.warmSnapshot || !this.warmTraceActive || this.warmSaved || !cache) return;
@@ -1156,6 +1169,7 @@ export class JVM {
       if (!trace) {
         this.fs.endWarmTrace();
         this.warmTraceActive = false;
+        this.dropJit();
         this.debug("warm snapshot busy: access record too large");
         return;
       }
@@ -1167,6 +1181,7 @@ export class JVM {
       this.fs.endWarmTrace();
       this.warmTraceActive = false;
       this.warmSaved = true;
+      this.dropJit();
       for (const c of this.created ?? []) c.bytes = undefined;
       this.created = null;
       this.debug("warm snapshot saved");
@@ -1262,6 +1277,7 @@ export class JVM {
           return false;
         }
         try {
+          this.dropJit();
           restoreVM(this, image, t);
           this.warmRestored = true;
           this.warmSaved = true;
@@ -1326,6 +1342,7 @@ export class JVM {
     if (!this.opts.warmSnapshot || !this.libraryStarting) {
       for (const c of this.created ?? []) c.bytes = undefined;
       this.created = null;
+      this.dropJit();
     }
     if (cmd !== undefined) {
       this.props.set("sun.java.command", cmd);

src/snapshot.ts

@@ -12,6 +12,7 @@ import { MAX_HEAP } from "./runtime.ts";
 import type { WarmTrace } from "./fs.ts";
 import { reopenSnapshotFile, OpenFile, streams } from "./natives/io.ts";
 import { JArray, JClass, JField, JMethod, JObject, Monitor } from "./runtime.ts";
+import { restoreJit } from "./jit.ts";
 import { Frame, JThread, NativeFrame, STATUS_RUNNABLE } from "./thread.ts";
 import type { NativeGen } from "./thread.ts";
 
@@ -48,6 +49,8 @@ interface ClassRec {
   nestHost: unknown;
   links: [string, unknown][];
   caches: [string, unknown][];
+  // generated JS and constants of compiled methods, in warm snapshots
+  jit?: [string, string, unknown][];
 }
 
 interface VmRec {
@@ -112,6 +115,18 @@ class Encoder {
 
   private work: (() => void)[] = [];
 
+  // enc, or null with nothing recorded when v cannot be saved
+  tryEnc(v: any, path: string): any {
+    const before = this.twins.size, pending = this.work.length;
+    try {
+      return this.enc(v, path);
+    } catch {
+      this.work.length = pending;
+      [...this.twins.keys()].slice(before).forEach((k) => this.twins.delete(k));
+      return null;
+    }
+  }
+
   // the twin of v; containers are filled in by drain, so deep object graphs
   // need no deep recursion
   enc(v: any, path: string): any {
@@ -271,6 +286,15 @@ export function dumpVM(vm: JVM, main: JThread, key: string, library = false): Sn
       links: c.methods.filter((m) => m.link).map((m) => [m.key, e.enc(m.link, `${p}.${m.key}.link`)]),
       caches: c.methods.filter((m) => m.cache).map((m) => [m.key, e.enc(m.cache, `${p}.${m.key}.cache`)]),
     };
+    if (library) {
+      const jit: [string, string, unknown][] = [];
+      for (const m of c.methods) {
+        const code = m.jit && vm.jitCode.get(m);
+        const consts = code && e.tryEnc(code.consts, `${p}.${m.key}.jit`);
+        if (code && consts) jit.push([m.key, code.src, consts]);
+      }
+      if (jit.length) rec.jit = jit;
+    }
     if (c.initThread || c.initWaiters.length) throw new Error(`snapshot: ${c.name} is being initialized`);
     if (rec.kind === "prim") rec.desc = c.prim!;
     else if (rec.kind === "array") rec.component = e.cls(c.component!, p);
@@ -397,6 +421,10 @@ export function restoreVM(vm: JVM, image: SnapshotImage, main: JThread) {
     c.nestHost = d.dec(r.nestHost);
     for (const [k, l] of r.links) c.methodMap.get(k)!.link = d.dec(l);
     for (const [k, cache] of r.caches ?? []) c.methodMap.get(k)!.cache = d.dec(cache);
+    for (const [k, src, consts] of r.jit ?? []) {
+      const m = c.methodMap.get(k);
+      if (m) restoreJit(vm, m, { src, consts: d.dec(consts) });
+    }
   });
   d.flush();
   const v = image.vm;

test/warm-jit.test.ts

@@ -0,0 +1,50 @@
+import { execFileSync } from "node:child_process";
+import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import { startLibrary } from "../src/node.ts";
+import { hotspot } from "../scripts/hotspot.ts";
+
+// Own file: a warm image restores reliably only in a process that has run no other VM.
+const root = join(import.meta.dirname, "..");
+const classes = join(root, "build", "test-warm-snapshot");
+mkdirSync(classes, { recursive: true });
+execFileSync(hotspot("javac"), ["-d", classes, join(root, "test", "warm-snapshot", "WarmJavac.java")]);
+
+async function compile(cache: string, args: string[]) {
+  const vm = await startLibrary({ classpath: [classes], snapshotDir: cache, warmSnapshot: true });
+  try {
+    const before = vm.jitCode.size;
+    const result = await vm.callStatic("WarmJavac", "compile", args);
+    return { result, before, after: vm.jitCode.size };
+  } finally {
+    vm.halt(0);
+  }
+}
+
+test("a warm restore compiles methods from their saved code on first call", async () => {
+  const dir = mkdtempSync(join(tmpdir(), "java-web-warm-jit-"));
+  try {
+    const src = join(dir, "Result.java");
+    writeFileSync(src, "public class Result { public static int value() { return 7; } }\n");
+    const reference = join(dir, "reference");
+    mkdirSync(reference);
+    execFileSync(hotspot("javac"), ["-g:none", "-d", reference, src]);
+    const expected = readFileSync(join(reference, "Result.class")).toString("base64");
+    const args = [src, join(dir, "not-created", "Absent.class"), join(dir, "out")];
+
+    const first = await compile(dir, args);
+    assert.equal(first.result, expected);
+    assert.ok(first.before < 100, "the first run has no saved code");
+
+    const restored = await compile(dir, args);
+    assert.equal(restored.result, expected);
+    assert.ok(restored.before > 100, `saved code for ${restored.before} methods`);
+    // compile removes a method's saved code only when it compiles from it
+    assert.ok(restored.before - restored.after >= 50, `${restored.before - restored.after} methods compiled from saved code`);
+  } finally {
+    rmSync(dir, { recursive: true, force: true });
+  }
+});