feat(core): support creating snapshots from existing snapshots (#14744)

2 files changed, 138 insertions, 25 deletions

diff --git a/core/bindings.rs b/core/bindings.rs
index b50c77df2..d7df613c5 100644
--- a/core/bindings.rs
+++ b/core/bindings.rs
@@ -7,6 +7,7 @@ use crate::modules::validate_import_assertions;
 use crate::modules::ImportAssertionsKind;
 use crate::modules::ModuleMap;
 use crate::ops::OpCtx;
+use crate::runtime::SnapshotOptions;
 use crate::JsRuntime;
 use log::debug;
 use std::option::Option;
@@ -96,8 +97,7 @@ pub fn module_origin<'a>(
 pub fn initialize_context<'s>(
   scope: &mut v8::HandleScope<'s, ()>,
   op_ctxs: &[OpCtx],
-  snapshot_loaded: bool,
-  will_snapshot: bool,
+  snapshot_options: SnapshotOptions,
 ) -> v8::Local<'s, v8::Context> {
   let scope = &mut v8::EscapableHandleScope::new(scope);
 
@@ -108,12 +108,12 @@ pub fn initialize_context<'s>(
 
   // Snapshot already registered `Deno.core.ops` but
   // extensions may provide ops that aren't part of the snapshot.
-  if snapshot_loaded {
+  if snapshot_options.loaded() {
     // Grab the Deno.core.ops object & init it
     let ops_obj = JsRuntime::grab_global::<v8::Object>(scope, "Deno.core.ops")
       .expect("Deno.core.ops to exist");
-    initialize_ops(scope, ops_obj, op_ctxs, snapshot_loaded);
-    if !will_snapshot {
+    initialize_ops(scope, ops_obj, op_ctxs, snapshot_options);
+    if snapshot_options != SnapshotOptions::CreateFromExisting {
       initialize_async_ops_info(scope, ops_obj, op_ctxs);
     }
     return scope.escape(context);
@@ -127,10 +127,10 @@ pub fn initialize_context<'s>(
 
   // Bind functions to Deno.core.ops.*
   let ops_obj = JsRuntime::ensure_objs(scope, global, "Deno.core.ops").unwrap();
-  if !will_snapshot {
+  if !snapshot_options.will_snapshot() {
     initialize_async_ops_info(scope, ops_obj, op_ctxs);
   }
-  initialize_ops(scope, ops_obj, op_ctxs, !will_snapshot);
+  initialize_ops(scope, ops_obj, op_ctxs, snapshot_options);
   scope.escape(context)
 }
 
@@ -138,14 +138,14 @@ fn initialize_ops(
   scope: &mut v8::HandleScope,
   ops_obj: v8::Local<v8::Object>,
   op_ctxs: &[OpCtx],
-  snapshot_loaded: bool,
+  snapshot_options: SnapshotOptions,
 ) {
   for ctx in op_ctxs {
     let ctx_ptr = ctx as *const OpCtx as *const c_void;
 
     // If this is a fast op, we don't want it to be in the snapshot.
     // Only initialize once snapshot is loaded.
-    if ctx.decl.fast_fn.is_some() && snapshot_loaded {
+    if ctx.decl.fast_fn.is_some() && snapshot_options.loaded() {
       set_func_raw(
         scope,
         ops_obj,
@@ -153,7 +153,7 @@ fn initialize_ops(
         ctx.decl.v8_fn_ptr,
         ctx_ptr,
         &ctx.decl.fast_fn,
-        snapshot_loaded,
+        snapshot_options,
       );
     } else {
       set_func_raw(
@@ -163,7 +163,7 @@ fn initialize_ops(
         ctx.decl.v8_fn_ptr,
         ctx_ptr,
         &None,
-        snapshot_loaded,
+        snapshot_options,
       );
     }
   }
@@ -190,7 +190,7 @@ pub fn set_func_raw(
   callback: v8::FunctionCallback,
   external_data: *const c_void,
   fast_function: &Option<Box<dyn FastFunction>>,
-  snapshot_loaded: bool,
+  snapshot_options: SnapshotOptions,
 ) {
   let key = v8::String::new(scope, name).unwrap();
   let external = v8::External::new(scope, external_data as *mut c_void);
@@ -198,11 +198,14 @@ pub fn set_func_raw(
     v8::FunctionTemplate::builder_raw(callback).data(external.into());
   let templ = if let Some(fast_function) = fast_function {
     // Don't initialize fast ops when snapshotting, the external references count mismatch.
-    if !snapshot_loaded {
-      builder.build(scope)
-    } else {
+    if matches!(
+      snapshot_options,
+      SnapshotOptions::Load | SnapshotOptions::None
+    ) {
       // TODO(@littledivy): Support fast api overloads in ops.
       builder.build_fast(scope, &**fast_function, None)
+    } else {
+      builder.build(scope)
     }
   } else {
     builder.build(scope)
diff --git a/core/runtime.rs b/core/runtime.rs
index e191596da..f27b3b337 100644
--- a/core/runtime.rs
+++ b/core/runtime.rs
@@ -82,6 +82,7 @@ pub struct JsRuntime {
   // This is an Option<OwnedIsolate> instead of just OwnedIsolate to workaround
   // a safety issue with SnapshotCreator. See JsRuntime::drop.
   v8_isolate: Option<v8::OwnedIsolate>,
+  snapshot_options: SnapshotOptions,
   built_from_snapshot: bool,
   allocations: IsolateAllocations,
   extensions: Vec<Extension>,
@@ -279,6 +280,38 @@ pub struct RuntimeOptions {
   pub inspector: bool,
 }
 
+#[derive(Copy, Clone, PartialEq, Eq)]
+pub enum SnapshotOptions {
+  Load,
+  CreateFromExisting,
+  Create,
+  None,
+}
+
+impl SnapshotOptions {
+  pub fn loaded(&self) -> bool {
+    matches!(
+      self,
+      SnapshotOptions::Load | SnapshotOptions::CreateFromExisting
+    )
+  }
+  pub fn will_snapshot(&self) -> bool {
+    matches!(
+      self,
+      SnapshotOptions::Create | SnapshotOptions::CreateFromExisting
+    )
+  }
+
+  fn from_bools(snapshot_loaded: bool, will_snapshot: bool) -> Self {
+    match (snapshot_loaded, will_snapshot) {
+      (true, true) => SnapshotOptions::CreateFromExisting,
+      (false, true) => SnapshotOptions::Create,
+      (true, false) => SnapshotOptions::Load,
+      (false, false) => SnapshotOptions::None,
+    }
+  }
+}
+
 impl Drop for JsRuntime {
   fn drop(&mut self) {
     if let Some(v8_isolate) = self.v8_isolate.as_mut() {
@@ -371,10 +404,39 @@ impl JsRuntime {
     let refs: &'static v8::ExternalReferences = Box::leak(Box::new(refs));
     let global_context;
 
-    let mut isolate = if options.will_snapshot {
-      // TODO(ry) Support loading snapshots before snapshotting.
-      assert!(options.startup_snapshot.is_none());
-      let snapshot_creator = v8::Isolate::snapshot_creator(Some(refs));
+    let (mut isolate, snapshot_options) = if options.will_snapshot {
+      let (snapshot_creator, snapshot_loaded) =
+        if let Some(snapshot) = options.startup_snapshot {
+          (
+            match snapshot {
+              Snapshot::Static(data) => {
+                v8::Isolate::snapshot_creator_from_existing_snapshot(
+                  data,
+                  Some(refs),
+                )
+              }
+              Snapshot::JustCreated(data) => {
+                v8::Isolate::snapshot_creator_from_existing_snapshot(
+                  data,
+                  Some(refs),
+                )
+              }
+              Snapshot::Boxed(data) => {
+                v8::Isolate::snapshot_creator_from_existing_snapshot(
+                  data,
+                  Some(refs),
+                )
+              }
+            },
+            true,
+          )
+        } else {
+          (v8::Isolate::snapshot_creator(Some(refs)), false)
+        };
+
+      let snapshot_options =
+        SnapshotOptions::from_bools(snapshot_loaded, options.will_snapshot);
+
       let mut isolate = JsRuntime::setup_isolate(snapshot_creator);
       {
         // SAFETY: this is first use of `isolate_ptr` so we are sure we're
@@ -385,11 +447,11 @@ impl JsRuntime {
         };
         let scope = &mut v8::HandleScope::new(&mut isolate);
         let context =
-          bindings::initialize_context(scope, &op_ctxs, false, true);
+          bindings::initialize_context(scope, &op_ctxs, snapshot_options);
         global_context = v8::Global::new(scope, context);
         scope.set_default_context(context);
       }
-      isolate
+      (isolate, snapshot_options)
     } else {
       let mut params = options
         .create_params
@@ -412,6 +474,9 @@ impl JsRuntime {
         false
       };
 
+      let snapshot_options =
+        SnapshotOptions::from_bools(snapshot_loaded, options.will_snapshot);
+
       let isolate = v8::Isolate::new(params);
       let mut isolate = JsRuntime::setup_isolate(isolate);
       {
@@ -423,12 +488,12 @@ impl JsRuntime {
         };
         let scope = &mut v8::HandleScope::new(&mut isolate);
         let context =
-          bindings::initialize_context(scope, &op_ctxs, snapshot_loaded, false);
+          bindings::initialize_context(scope, &op_ctxs, snapshot_options);
 
         global_context = v8::Global::new(scope, context);
       }
 
-      isolate
+      (isolate, snapshot_options)
     };
 
     op_state.borrow_mut().put(isolate_ptr);
@@ -464,6 +529,7 @@ impl JsRuntime {
     let mut js_runtime = Self {
       v8_isolate: Some(isolate),
       built_from_snapshot: has_startup_snapshot,
+      snapshot_options,
       allocations: IsolateAllocations::default(),
       event_loop_middlewares: Vec::with_capacity(options.extensions.len()),
       extensions: options.extensions,
@@ -550,8 +616,10 @@ impl JsRuntime {
       let context = bindings::initialize_context(
         scope,
         &self.state.borrow().op_ctxs,
-        self.built_from_snapshot,
-        false,
+        SnapshotOptions::from_bools(
+          self.built_from_snapshot,
+          self.snapshot_options.will_snapshot(),
+        ),
       );
       JsRealm::new(v8::Global::new(scope, context))
     };
@@ -2749,6 +2817,48 @@ pub mod tests {
       .execute_script("check.js", "if (a != 3) throw Error('x')")
       .unwrap();
   }
+
+  #[test]
+  fn will_snapshot2() {
+    let startup_data = {
+      let mut runtime = JsRuntime::new(RuntimeOptions {
+        will_snapshot: true,
+        ..Default::default()
+      });
+      runtime.execute_script("a.js", "let a = 1 + 2").unwrap();
+      runtime.snapshot()
+    };
+
+    let snapshot = Snapshot::JustCreated(startup_data);
+    let mut runtime = JsRuntime::new(RuntimeOptions {
+      will_snapshot: true,
+      startup_snapshot: Some(snapshot),
+      ..Default::default()
+    });
+
+    let startup_data = {
+      runtime
+        .execute_script("check_a.js", "if (a != 3) throw Error('x')")
+        .unwrap();
+      runtime.execute_script("b.js", "b = 2 + 3").unwrap();
+      runtime.snapshot()
+    };
+
+    let snapshot = Snapshot::JustCreated(startup_data);
+    {
+      let mut runtime = JsRuntime::new(RuntimeOptions {
+        startup_snapshot: Some(snapshot),
+        ..Default::default()
+      });
+      runtime
+        .execute_script("check_b.js", "if (b != 5) throw Error('x')")
+        .unwrap();
+      runtime
+        .execute_script("check2.js", "if (!Deno.core) throw Error('x')")
+        .unwrap();
+    }
+  }
+
   #[test]
   fn test_snapshot_callbacks() {
     let snapshot = {