Experimental code cache for built-in js/ts modules

src/workerd/jsg/modules.c++

@@ -4,10 +4,40 @@
 
 #include "modules.h"
 #include "promise.h"
+#include <kj/mutex.h>
 
 namespace workerd::jsg {
 namespace {
 
+class CompileCache {
+  // The CompileCache is used to hold cached compilation data for built-in JavaScript modules.
+  //
+  // Importantly, this is a process-lifetime in-memory cache that is only appropriate for
+  // built-in modules.
+  //
+  // The memory-safety of this cache depends on the assumption that entries are never removed
+  // or replaced. If things are ever changed such that entries are removed/replaced, then
+  // we'd likely need to have find return an atomic refcount or something similar.
+public:
+  void add(const void* key, std::unique_ptr<v8::ScriptCompiler::CachedData> cached) const {
+    cache.lockExclusive()->upsert(key, kj::mv(cached), [](auto&,auto&&) {});
+  }
+
+  kj::Maybe<v8::ScriptCompiler::CachedData&> find(const void* key) const {
+    return cache.lockShared()->find(key).map([](auto& data)
+        -> v8::ScriptCompiler::CachedData& { return *data; });
+  }
+
+  static const CompileCache& get() {
+    static CompileCache instance;
+    return instance;
+  }
+
+private:
+  kj::MutexGuarded<kj::HashMap<const void*, std::unique_ptr<v8::ScriptCompiler::CachedData>>> cache;
+  // The key is the address of the static global that was compiled to produce the CachedData.
+};
+
 ModuleRegistry* getModulesForResolveCallback(v8::Isolate* isolate) {
   return static_cast<ModuleRegistry*>(
       isolate->GetCurrentContext()->GetAlignedPointerFromEmbedderData(2));
@@ -275,7 +305,7 @@ v8::Local<v8::Module> compileEsmModule(
     jsg::Lock& js,
     kj::StringPtr name,
     kj::ArrayPtr<const char> content,
-    ModuleInfoCompileFlags flags) {
+    ModuleInfoCompileOption option) {
   // Must pass true for `is_module`, but we can skip everything else.
   const int resourceLineOffset = 0;
   const int resourceColumnOffset = 0;
@@ -291,16 +321,39 @@ v8::Local<v8::Module> compileEsmModule(
                           resourceIsSharedCrossOrigin, scriptId, {},
                           resourceIsOpaque, isWasm, isModule);
   v8::Local<v8::String> contentStr;
-  if ((flags & ModuleInfoCompileFlags::EXTERNAL) == ModuleInfoCompileFlags::EXTERNAL) {
+
+  if (option == ModuleInfoCompileOption::BUILTIN) {
     // TODO(later): Use of newExternalOneByteString here limits our built-in source
     // modules (for which this path is used) to only the latin1 character set. We
     // may need to revisit that to import built-ins as UTF-16 (two-byte).
     contentStr = jsg::check(jsg::newExternalOneByteString(js, content));
-  } else {
-    contentStr = jsg::v8Str(js.v8Isolate, content);
+
+    const auto& compileCache = CompileCache::get();
+    KJ_IF_MAYBE(cached, compileCache.find(content.begin())) {
+      v8::ScriptCompiler::Source source(contentStr, origin, cached);
+      v8::ScriptCompiler::CompileOptions options = v8::ScriptCompiler::kConsumeCodeCache;
+      KJ_DEFER(if (source.GetCachedData()->rejected) {
+        KJ_LOG(ERROR, kj::str("Failed to load module '", name ,"' using compile cache"));
+        js.throwException(KJ_EXCEPTION(FAILED, "jsg.Error: Internal error"));
+      });
+      return jsg::check(v8::ScriptCompiler::CompileModule(js.v8Isolate, &source, options));
+    }
+
+    v8::ScriptCompiler::Source source(contentStr, origin);
+    auto module = jsg::check(v8::ScriptCompiler::CompileModule(js.v8Isolate, &source));
+
+    auto cachedData = std::unique_ptr<v8::ScriptCompiler::CachedData>(
+        v8::ScriptCompiler::CreateCodeCache(module->GetUnboundModuleScript()));
+    compileCache.add(content.begin(), kj::mv(cachedData));
+    return module;
   }
+
+  contentStr = jsg::v8Str(js.v8Isolate, content);
+
   v8::ScriptCompiler::Source source(contentStr, origin);
-  return jsg::check(v8::ScriptCompiler::CompileModule(js.v8Isolate, &source));
+  auto module = jsg::check(v8::ScriptCompiler::CompileModule(js.v8Isolate, &source));
+
+  return module;
 }
 
 v8::Local<v8::Module> createSyntheticModule(
@@ -333,7 +386,7 @@ ModuleRegistry::ModuleInfo::ModuleInfo(
     jsg::Lock& js,
     kj::StringPtr name,
     kj::ArrayPtr<const char> content,
-    CompileFlags flags)
+    CompileOption flags)
     : ModuleInfo(js, compileEsmModule(js, name, content, flags)) {}
 
 ModuleRegistry::ModuleInfo::ModuleInfo(

src/workerd/jsg/modules.h

@@ -175,21 +175,21 @@ class ModuleRegistry {
                v8::Local<v8::Module> module,
                kj::Maybe<SyntheticModuleInfo> maybeSynthetic = nullptr);
 
-    enum class CompileFlags {
-      NONE,
-      EXTERNAL,
-      // The EXTERNAL flag tells the compile operation to treat the content as an external
-      // string wrapping an immutable buffer outside the V8 heap. When this flag is not set,
-      // the content is copied into the V8 heap.
-      // TODO(cleanup): Once we have a more complete set of options here for options that
-      // are used for all built-in modules, we'll likely collapse the flags into a single
-      // BUILTIN flag that encompasses multiple options.
+    enum class CompileOption {
+      BUNDLE,
+      // The BUNDLE options tells the compile operation to threat the content as coming
+      // from a worker bundle.
+      BUILTIN,
+      // The BUILTIN option tells the compile operation to treat the content as a builtin
+      // module. This implies certain changes in behavior, such as treating the content
+      // as an immutable, process-lifetime buffer that will never be destroyed, and caching
+      // the compilation data.
     };
 
     ModuleInfo(jsg::Lock& js,
                kj::StringPtr name,
                kj::ArrayPtr<const char> content,
-               CompileFlags flags = CompileFlags::NONE);
+               CompileOption flags = CompileOption::BUNDLE);
 
     ModuleInfo(jsg::Lock& js, kj::StringPtr name,
                kj::Maybe<kj::ArrayPtr<kj::StringPtr>> maybeExports,
@@ -225,18 +225,7 @@ class ModuleRegistry {
   virtual void setDynamicImportCallback(kj::Function<DynamicImportCallback> func) = 0;
 };
 
-using ModuleInfoCompileFlags = ModuleRegistry::ModuleInfo::CompileFlags;
-
-inline constexpr ModuleInfoCompileFlags operator|(
-    ModuleInfoCompileFlags a,
-    ModuleInfoCompileFlags b) {
-  return static_cast<ModuleInfoCompileFlags>(static_cast<uint>(a) | static_cast<uint>(b));
-}
-inline constexpr ModuleInfoCompileFlags operator&(
-    ModuleInfoCompileFlags a,
-    ModuleInfoCompileFlags b) {
-  return static_cast<ModuleInfoCompileFlags>(static_cast<uint>(a) & static_cast<uint>(b));
-}
+using ModuleInfoCompileOption = ModuleRegistry::ModuleInfo::CompileOption;
 
 template <typename TypeWrapper>
 class ModuleRegistryImpl final: public ModuleRegistry {
@@ -429,7 +418,7 @@ class ModuleRegistryImpl final: public ModuleRegistry {
           return moduleInfo;
         }
         KJ_CASE_ONEOF(src, kj::ArrayPtr<const char>) {
-          info = ModuleInfo(js, specifier.toString(), src, ModuleInfoCompileFlags::EXTERNAL);
+          info = ModuleInfo(js, specifier.toString(), src, ModuleInfoCompileOption::BUILTIN);
           return KJ_ASSERT_NONNULL(info.tryGet<ModuleInfo>());
         }
         KJ_CASE_ONEOF(src, kj::Function<ModuleInfo(Lock&)>) {