Merge pull request #3293 from kevinzakka:studio-live-parallel-prefetch

PiperOrigin-RevId: 922417778
Change-Id: Ia839747c4af0e705962308dc84e9a7c3926fc910

src/experimental/studio/emscripten.cc

@@ -27,6 +27,7 @@
 #include <string_view>
 #include <unordered_map>
 #include <utility>
+#include <vector>
 
 #include <mujoco/mujoco.h>
 #include "experimental/platform/hal/graphics_mode.h"
@@ -108,8 +109,9 @@ EM_ASYNC_JS(int, FetchUrl,
               }
             });
 
-// Cache for data fetched via HTTP/HTTPS. Stores the downloaded bytes keyed by
+// Cache for data fetched via HTTP/HTTPS. Stores downloaded (or pre-primed)
-// the resource name (URL) so that read() can return a pointer to the data.
+// bytes keyed by the resource URL so that the resource provider's read()
+// can return a pointer into stable storage.
 class FetchCache {
  public:
   static FetchCache& Instance() {
@@ -117,52 +119,74 @@ class FetchCache {
     return instance;
   }
 
-  // Fetches the URL and stores the result. Returns the size (>0) on success.
+  // Returns the size of the entry for `url`, fetching it over the network
+  // if necessary. Entries that have been pre-populated via Prime() are
+  // returned without touching the network. Returns 0 on fetch failure.
   int Fetch(const char* url) {
-    char* data = nullptr;
+    if (auto it = entries_.find(url); it != entries_.end()) {
+      return static_cast<int>(it->second.size());
+    }
+    char* buf = nullptr;
     std::int32_t size = 0;
-    if (!FetchUrl(url, &data, &size)) {
+    if (!FetchUrl(url, &buf, &size)) {
       return 0;
     }
-    entries_[url] = Entry{UniquePtrWasm<char[]>(data), size};
+    // FetchUrl hands us a malloc'd buffer; take ownership of the bytes via
-    return size;
+    // a std::string and release the original allocation.
+    std::string bytes(buf, size);
+    std::free(buf);
+    return static_cast<int>(entries_.emplace(url, std::move(bytes))
+                                .first->second.size());
   }
 
-  // Returns pointer and size for a previously fetched URL.
+  // Pre-populates the cache so a subsequent Fetch() of the same URL hits
-  int Read(const char* url, const void** buffer) {
+  // memory instead of the network. First writer wins.
+  void Prime(std::string url, std::string bytes) {
+    entries_.try_emplace(std::move(url), std::move(bytes));
+  }
+
+  // Hands back a pointer into the cached bytes for the given URL.
+  // Returns -1 if the URL has not been fetched or primed.
+  int Read(const char* url, const void** buffer) const {
     auto it = entries_.find(url);
-    if (it == entries_.end()) {
+    if (it == entries_.end()) return -1;
-      return -1;
+    *buffer = it->second.data();
-    }
+    return static_cast<int>(it->second.size());
-    *buffer = it->second.data.get();
-    return it->second.size;
   }
 
-  // Frees the data for a URL.
   void Close(const char* url) { entries_.erase(url); }
 
  private:
-  struct FreeDeleter {
+  std::unordered_map<std::string, std::string> entries_;
-    void operator()(void* p) const { std::free(p); }
-  };
-  template <typename T>
-  using UniquePtrWasm = std::unique_ptr<T, FreeDeleter>;
-
-  struct Entry {
-    UniquePtrWasm<char[]> data;
-    int size;
-  };
-  std::unordered_map<std::string, Entry> entries_;
 };
 
 // ---------------------------------------------------------------------------
 
-// Javascript-facing function to register an asset.
+// Javascript-facing function to register a build-time asset (font, IBL,
+// .filamat material) that is fetched via the static AssetRegistry rather
+// than over HTTP.
 void RegisterAsset(std::string filename, std::string contents) {
   AssetRegistry::Instance().RegisterAsset(std::move(filename),
                                           std::move(contents));
 }
 
+// Javascript-facing function to pre-populate the HTTP/HTTPS fetch cache.
+// Used by the JS-side parallel prefetcher to feed already-downloaded asset
+// bytes to MuJoCo's compiler before LoadUrl is invoked, sidestepping the
+// otherwise strictly-serial EM_ASYNC_JS fetch chain.
+void PrimeFetchCache(std::string url, std::string bytes) {
+  FetchCache::Instance().Prime(std::move(url), std::move(bytes));
+}
+
+// Javascript-facing wrapper around mju_getXMLDependencies. Returns the
+// list of resource URLs that the given root XML transitively references
+// (meshes, textures, includes, ...), per MuJoCo's compiler rules.
+std::vector<std::string> GetXMLDependencies(const std::string& root_url) {
+  mjStringVec deps;
+  mju_getXMLDependencies(root_url.c_str(), &deps);
+  return deps;
+}
+
 // Javascript-facing function to initialize the app.
 void Init(const std::string& title, bool dark_theme) {
   // Note: dimensions do not matter as window will be resized to fit canvas.
@@ -305,7 +329,11 @@ void Deinit() {
 }
 
 EMSCRIPTEN_BINDINGS(studio_bindings) {
+  // Matches the C-level `mjStringVec` typedef in mjspec.h.
+  emscripten::register_vector<std::string>("mjStringVec");
   emscripten::function("registerAsset", &RegisterAsset);
+  emscripten::function("primeFetchCache", &PrimeFetchCache);
+  emscripten::function("getXMLDependencies", &GetXMLDependencies);
   emscripten::function("init", &Init);
   emscripten::function("loadFile", &LoadFile);
   emscripten::function("loadUrl", &LoadUrl);

src/experimental/studio/live.js

@@ -21,6 +21,85 @@ function hideLoading() {
   loadingOverlay.style.display = 'none';
 }
 
+// ---------------------------------------------------------------------------
+// Parallel asset prefetcher.
+//
+// MuJoCo's XML compiler synchronously opens each referenced asset via
+// mjpResourceProvider. In WASM that becomes a chain of ASYNCIFY-suspended
+// fetch() calls, strictly serial, one RTT per asset. For Menagerie models
+// that's ~80 round trips and ~10 s of avoidable wall time on a cold cache.
+//
+// To sidestep this we let MuJoCo itself enumerate the dependencies (via
+// mju_getXMLDependencies, exposed to JS as Module.getXMLDependencies), then
+// fetch them all in parallel and feed the bytes into the WASM-side
+// FetchCache via Module.primeFetchCache. When Module.loadUrl subsequently
+// runs, every resource the compiler opens is already in memory.
+//
+// Discovery rules live in src/xml/xml_util.cc; the JS only has to know
+// the scheme rewrites used by the resource providers in emscripten.cc.
+// ---------------------------------------------------------------------------
+
+// MuJoCo Live custom URL schemes that the C++ resource providers rewrite
+// before fetching. JS must do the same so primed cache keys match what
+// the providers eventually request.
+function resolveScheme(url) {
+  if (url.startsWith('github:')) {
+    return 'https://raw.githubusercontent.com/' + url.slice('github:'.length);
+  }
+  return url;
+}
+
+async function prefetchModelAssets(rootUrl, onProgress) {
+  const primed = new Set();   // URLs we've already pushed into FetchCache
+  const stats = { files: 0, bytes: 0, errors: 0 };
+
+  // Fetch + prime, dedup against `primed`. Returns true on success.
+  async function fetchAndPrime(httpUrl) {
+    if (primed.has(httpUrl)) return true;
+    primed.add(httpUrl);
+    try {
+      const resp = await fetch(httpUrl);
+      if (!resp.ok) {
+        console.warn('[prefetch] HTTP', resp.status, httpUrl);
+        stats.errors++;
+        return false;
+      }
+      const bytes = new Uint8Array(await resp.arrayBuffer());
+      Module.primeFetchCache(httpUrl, bytes);
+      stats.files++;
+      stats.bytes += bytes.length;
+      onProgress?.(stats);
+      return true;
+    } catch (err) {
+      console.warn('[prefetch]', httpUrl, err);
+      stats.errors++;
+      return false;
+    }
+  }
+
+  // Step 1: fetch the root XML so mju_getXMLDependencies (which reads it
+  // through the resource provider) hits cache instead of the network.
+  // Also gives the loading screen an immediate progress tick.
+  const rootHttp = resolveScheme(rootUrl);
+  if (!await fetchAndPrime(rootHttp)) {
+    throw new Error(`could not fetch root model ${rootHttp}`);
+  }
+
+  // Step 2: ask MuJoCo for the transitive dependency list. The call is
+  // async because reading recursively included XMLs may suspend on
+  // ASYNCIFY fetches; those XMLs end up in FetchCache as a side effect.
+  // The result is an embind mjStringVec we copy out and release.
+  const depsVec = await Module.getXMLDependencies(rootUrl);
+  const deps = [];
+  for (let i = 0; i < depsVec.size(); ++i) deps.push(depsVec.get(i));
+  depsVec.delete();
+
+  // Step 3: fetch every dependency in parallel, deduped against URLs
+  // already primed in step 1 (and earlier deps within this batch).
+  await Promise.all(deps.map((depUrl) => fetchAndPrime(resolveScheme(depUrl))));
+  return stats;
+}
+
 var Module = {
   preRun: [],
   postRun: [],
@@ -111,13 +190,25 @@ var Module = {
             // loop until the load completes, otherwise renderFrame will
             // hit "Cannot have multiple async operations in flight".
             showLoading();
+            // Surface prefetch progress on the loading screen.
+            const dlEl = document.getElementById('loadingDownload');
+            const dlFileEl = document.getElementById('loadingDownloadFile');
+            if (dlEl) dlEl.style.display = 'block';
+            const onProgress = (s) => {
+              if (dlFileEl) {
+                dlFileEl.textContent =
+                    `${s.files} files · ${(s.bytes/1024/1024).toFixed(1)} MB`;
+              }
+            };
             requestAnimationFrame(() => {
               requestAnimationFrame(async () => {
                 try {
+                  await prefetchModelAssets(modelUrl, onProgress);
                   await Module.loadUrl(modelUrl);
                 } catch (error) {
                   console.error('Failed to load model from URL:', error);
                 } finally {
+                  if (dlEl) dlEl.style.display = 'none';
                   hideLoading();
                   requestAnimationFrame(Module.animate);
                 }

src/xml/xml_util.cc

@@ -222,34 +222,46 @@ void mjCopyError(char* dst, const char* src, int maxlen) {
 }
 
 void mju_getXMLDependencies(const char* filename, mjStringVec* dependencies) {
-  // load XML file or parse string
+  // Open through the resource provider so the function works against any
-  tinyxml2::XMLDocument doc;
+  // registered backend (OS file system, VFS, HTTP, "github:", ...) rather
-  doc.LoadFile(filename);
+  // than only the OS file system.
-
+  mjResource* resource = mju_openResource("", filename, nullptr, nullptr, 0);
-  // error checking
+  if (resource == nullptr) {
-  if (doc.Error()) {
+    mju_error("Could not open '%s'", filename);
-    mju_error("Problem reading XML file '%s': %s", filename, doc.ErrorStr());
   }
 
-  // get top-level element
+  // Read the XML bytes from the resource.
-  tinyxml2::XMLElement *root = doc.RootElement();
+  const void* buffer = nullptr;
-  if (!root) {
+  int size = mju_readResource(resource, &buffer);
-    mju_error("XML root element not found");
+  if (size < 0 || !size) {
-  }
-  std::unordered_set<std::string> files = {filename};
-
-  std::optional<FilePath> model_dir = std::nullopt;
-  mjResource *resource = mju_openResource("", filename, nullptr,
-                                          nullptr, 0);
-  if (resource != nullptr) {
-    const char* dir;
-    int ndir;
-    mju_getResourceDir(resource, &dir, &ndir);
-    model_dir = FilePath(std::string(dir, ndir));
     mju_closeResource(resource);
+    mju_error("Could not read '%s'", filename);
   }
-  // Get file references from include and model tags.
+
-  AccumulateFiles(files, root, model_dir.value());
+  // Capture the model directory while the resource is still open.
+  const char* dir = nullptr;
+  int ndir = 0;
+  mju_getResourceDir(resource, &dir, &ndir);
+  FilePath model_dir(std::string(dir, ndir));
+
+  // Parse from buffer and close (the parsed DOM is independent of the
+  // resource buffer once Parse returns).
+  tinyxml2::XMLDocument doc;
+  tinyxml2::XMLError err =
+      doc.Parse(static_cast<const char*>(buffer), static_cast<size_t>(size));
+  mju_closeResource(resource);
+
+  if (err != tinyxml2::XML_SUCCESS) {
+    mju_error("Problem reading XML file '%s': %s", filename,
+              doc.ErrorStr() ? doc.ErrorStr() : "");
+  }
+  tinyxml2::XMLElement* root = doc.RootElement();
+  if (!root) {
+    mju_error("XML root element not found in '%s'", filename);
+  }
+
+  std::unordered_set<std::string> files = {filename};
+  AccumulateFiles(files, root, model_dir);
 
   *dependencies = {files.begin(), files.end()};
 }

test/xml/xml_utils_test.cc

@@ -13,11 +13,14 @@
 // limitations under the License.
 
 
+#include <map>
 #include <set>
 #include <string>
 #include <vector>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include <mujoco/mjplugin.h>
 #include <mujoco/mujoco.h>
 #include "test/fixture.h"
 
@@ -42,5 +45,64 @@ TEST_F(MujocoTest, GetXMLDependenciesTest) {
                                                 kModelPaths.end()};
   EXPECT_EQ(dependency_set, expected_dependency_set);
 }
+
+// Custom resource provider that serves XML strings from an in-memory map,
+// used to exercise the non-filesystem code path.
+namespace memxml {
+static const std::map<std::string, std::string>* g_files = nullptr;
+
+int Open(mjResource* resource) {
+  return g_files && g_files->count(resource->name) ? 1 : 0;
+}
+int Read(mjResource* resource, const void** buffer) {
+  auto it = g_files->find(resource->name);
+  if (it == g_files->end()) return -1;
+  *buffer = it->second.data();
+  return static_cast<int>(it->second.size());
+}
+void Close(mjResource* resource) {}
+}  // namespace memxml
+
+// Verifies mju_getXMLDependencies works against a non-filesystem resource
+// provider (the case the WASM/HTTP build relies on).
+TEST_F(MujocoTest, GetXMLDependenciesViaResourceProvider) {
+  const std::map<std::string, std::string> files = {
+      {"memxml:/scene.xml",
+       "<mujoco>\n"
+       "  <include file=\"child.xml\"/>\n"
+       "  <compiler meshdir=\"meshes\"/>\n"
+       "  <asset>\n"
+       "    <mesh name=\"m\" file=\"m.obj\"/>\n"
+       "  </asset>\n"
+       "</mujoco>"},
+      {"memxml:/child.xml",
+       "<mujoco>\n"
+       "  <asset>\n"
+       "    <texture name=\"t\" type=\"2d\" file=\"t.png\"/>\n"
+       "  </asset>\n"
+       "</mujoco>"},
+  };
+  memxml::g_files = &files;
+
+  mjpResourceProvider provider = {
+      .prefix = "memxml",
+      .open = memxml::Open,
+      .read = memxml::Read,
+      .close = memxml::Close,
+  };
+  ASSERT_GT(mjp_registerResourceProvider(&provider), 0);
+
+  mjStringVec dependencies;
+  mju_getXMLDependencies("memxml:/scene.xml", &dependencies);
+  std::set<std::string> dep_set(dependencies.begin(), dependencies.end());
+
+  EXPECT_THAT(dep_set, testing::UnorderedElementsAre(
+      "memxml:/scene.xml",
+      "memxml:/child.xml",
+      "memxml:/meshes/m.obj",
+      "memxml:/t.png"));
+
+  memxml::g_files = nullptr;
+}
 }  // namespace
 }  // namespace mujoco