Bug 1779347 - Support ESM in Cu.getModuleImportStack. r=jonco

Replaced ModuleEntry.importStack field with mozJSModuleLoader.mImportStacks map,
to collect the import stack both for JSM and ESM when it's enabled.

Differential Revision: https://phabricator.services.mozilla.com/D151684

browser/base/content/test/performance/StartupContentSubframe.jsm

@@ -41,7 +41,7 @@ class StartupContentSubframeChild extends JSWindowActorChild {
       modules[module] = collectStacks ? Cu.getModuleImportStack(module) : "";
     }
     for (let module of Cu.loadedESModules) {
-      modules[module] = "";
+      modules[module] = collectStacks ? Cu.getModuleImportStack(module) : "";
     }
 
     let services = {};

browser/base/content/test/performance/browser_startup_content.js

@@ -137,7 +137,9 @@ add_task(async function() {
             : "";
         }
         for (let module of Cu.loadedESModules) {
-          modules[module] = "";
+          modules[module] = collectStacks
+            ? Cu.getModuleImportStack(module)
+            : "";
         }
         let services = {};
         for (let contractID of Object.keys(Cc)) {

js/xpconnect/loader/mozJSModuleLoader.cpp

@@ -265,6 +265,9 @@ mozJSModuleLoader::mozJSModuleLoader()
     : mImports(16),
       mInProgressImports(16),
       mFallbackImports(16),
+#ifdef STARTUP_RECORDER_ENABLED
+      mImportStacks(16),
+#endif
       mLocations(16),
       mInitialized(false),
       mLoaderGlobal(dom::RootingCx()),
@@ -464,12 +467,29 @@ static size_t SizeOfTableExcludingThis(
   return n;
 }
 
+#ifdef STARTUP_RECORDER_ENABLED
+template <class Key, class Data, class UserData, class Converter>
+static size_t SizeOfStringTableExcludingThis(
+    const nsBaseHashtable<Key, Data, UserData, Converter>& aTable,
+    MallocSizeOf aMallocSizeOf) {
+  size_t n = aTable.ShallowSizeOfExcludingThis(aMallocSizeOf);
+  for (const auto& entry : aTable) {
+    n += entry.GetKey().SizeOfExcludingThisIfUnshared(aMallocSizeOf);
+    n += entry.GetData().SizeOfExcludingThisIfUnshared(aMallocSizeOf);
+  }
+  return n;
+}
+#endif
+
 size_t mozJSModuleLoader::SizeOfIncludingThis(MallocSizeOf aMallocSizeOf) {
   size_t n = aMallocSizeOf(this);
   n += SizeOfTableExcludingThis(mImports, aMallocSizeOf);
   n += mLocations.ShallowSizeOfExcludingThis(aMallocSizeOf);
   n += SizeOfTableExcludingThis(mInProgressImports, aMallocSizeOf);
   n += SizeOfTableExcludingThis(mFallbackImports, aMallocSizeOf);
+#ifdef STARTUP_RECORDER_ENABLED
+  n += SizeOfStringTableExcludingThis(mImportStacks, aMallocSizeOf);
+#endif
   return n;
 }
 
@@ -979,6 +999,9 @@ void mozJSModuleLoader::UnloadModules() {
   }
   mServicesObj = nullptr;
 
+#ifdef STARTUP_RECORDER_ENABLED
+  mImportStacks.Clear();
+#endif
   mFallbackImports.Clear();
   mInProgressImports.Clear();
   mImports.Clear();
@@ -1169,6 +1192,18 @@ nsresult mozJSModuleLoader::GetLoadedJSAndESModules(
   return NS_OK;
 }
 
+#ifdef STARTUP_RECORDER_ENABLED
+void mozJSModuleLoader::RecordImportStack(JSContext* aCx,
+                                          const nsACString& aLocation) {
+  if (!Preferences::GetBool("browser.startup.record", false)) {
+    return;
+  }
+
+  mImportStacks.InsertOrUpdate(
+      aLocation, xpc_PrintJSStack(aCx, false, false, false).get());
+}
+#endif
+
 nsresult mozJSModuleLoader::GetModuleImportStack(const nsACString& aLocation,
                                                  nsACString& retval) {
 #ifdef STARTUP_RECORDER_ENABLED
@@ -1176,13 +1211,12 @@ nsresult mozJSModuleLoader::GetModuleImportStack(const nsACString& aLocation,
   MOZ_ASSERT(mInitialized);
 
   ModuleLoaderInfo info(aLocation);
-
-  ModuleEntry* mod;
-  if (!mImports.Get(info.Key(), &mod)) {
+  auto str = mImportStacks.Lookup(info.Key());
+  if (!str) {
     return NS_ERROR_FAILURE;
   }
 
-  retval = mod->importStack;
+  retval = *str;
   return NS_OK;
 #else
   return NS_ERROR_NOT_IMPLEMENTED;
@@ -1441,9 +1475,7 @@ nsresult mozJSModuleLoader::Import(JSContext* aCx, const nsACString& aLocation,
     }
 
 #ifdef STARTUP_RECORDER_ENABLED
-    if (Preferences::GetBool("browser.startup.record", false)) {
-      newEntry->importStack = xpc_PrintJSStack(aCx, false, false, false).get();
-    }
+    RecordImportStack(aCx, aLocation);
 #endif
 
     mod = newEntry.get();
@@ -1590,6 +1622,12 @@ nsresult mozJSModuleLoader::ImportESModule(
   nsresult rv = NS_NewURI(getter_AddRefs(uri), aLocation);
   NS_ENSURE_SUCCESS(rv, rv);
 
+#ifdef STARTUP_RECORDER_ENABLED
+  if (!mModuleLoader->IsModuleFetched(uri)) {
+    RecordImportStack(aCx, aLocation);
+  }
+#endif
+
   nsCOMPtr<nsIPrincipal> principal =
       mModuleLoader->GetGlobalObject()->PrincipalOrNull();
   MOZ_ASSERT(principal);

js/xpconnect/loader/mozJSModuleLoader.h

@@ -91,6 +91,10 @@ class mozJSModuleLoader final : public nsIMemoryReporter {
       JS::MutableHandleObject aModuleExports, bool aIgnoreExports,
       bool* aFound);
 
+#ifdef STARTUP_RECORDER_ENABLED
+  void RecordImportStack(JSContext* aCx, const nsACString& aLocation);
+#endif
+
   nsresult Unload(const nsACString& aResourceURI);
   nsresult IsModuleLoaded(const nsACString& aResourceURI, bool* aRetval);
   nsresult IsJSModuleLoaded(const nsACString& aResourceURI, bool* aRetval);
@@ -181,9 +185,6 @@ class mozJSModuleLoader final : public nsIMemoryReporter {
       obj = nullptr;
       thisObjectKey = nullptr;
       location = nullptr;
-#ifdef STARTUP_RECORDER_ENABLED
-      importStack.Truncate();
-#endif
     }
 
     size_t SizeOfIncludingThis(mozilla::MallocSizeOf aMallocSizeOf) const;
@@ -193,9 +194,6 @@ class mozJSModuleLoader final : public nsIMemoryReporter {
     JS::PersistentRootedScript thisObjectKey;
     char* location;
     nsCString resolvedURL;
-#ifdef STARTUP_RECORDER_ENABLED
-    nsCString importStack;
-#endif
   };
 
   class FallbackModuleEntry {
@@ -224,6 +222,9 @@ class mozJSModuleLoader final : public nsIMemoryReporter {
   nsClassHashtable<nsCStringHashKey, ModuleEntry> mImports;
   nsTHashMap<nsCStringHashKey, ModuleEntry*> mInProgressImports;
   nsClassHashtable<nsCStringHashKey, FallbackModuleEntry> mFallbackImports;
+#ifdef STARTUP_RECORDER_ENABLED
+  nsTHashMap<nsCStringHashKey, nsCString> mImportStacks;
+#endif
 
   // A map of on-disk file locations which are loaded as modules to the
   // pre-resolved URIs they were loaded from. Used to prevent the same file

js/xpconnect/tests/unit/import_stack.jsm

@@ -0,0 +1,2 @@
+function test() {}
+var EXPORTED_SYMBOLS = ["test"];

js/xpconnect/tests/unit/import_stack.sys.mjs

@@ -0,0 +1 @@
+export function test() {}

js/xpconnect/tests/unit/test_import_stack.js

@@ -0,0 +1,16 @@
+Services.prefs.setBoolPref("browser.startup.record", true);
+registerCleanupFunction(() => {
+  Services.prefs.clearUserPref("browser.startup.record");
+});
+
+add_task(function test_JSModule() {
+  const URL = "resource://test/import_stack.jsm";
+  ChromeUtils.import(URL);
+  Assert.ok(Cu.getModuleImportStack(URL).includes("test_JSModule"));
+});
+
+add_task(function test_ESModule() {
+  const URL = "resource://test/import_stack.sys.mjs";
+  ChromeUtils.importESModule(URL);
+  Assert.ok(Cu.getModuleImportStack(URL).includes("test_ESModule"));
+});

js/xpconnect/tests/unit/xpcshell.ini

@@ -45,6 +45,8 @@ support-files =
   es6module_loaded-2.sys.mjs
   es6module_loaded-3.sys.mjs
   api_script.js
+  import_stack.jsm
+  import_stack.sys.mjs
 
 [test_allowWaivers.js]
 [test_bogus_files.js]
@@ -192,3 +194,7 @@ head = head_watchdog.js
 [test_defineESModuleGetters.js]
 [test_loadedESModules.js]
 [test_import_from_sandbox.js]
+[test_import_stack.js]
+skip-if =
+    !nightly_build
+    !debug