From f7fb85d60a5839b218916d8d54a331f390527716 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kat=20March=C3=A1n?= <kzm@zkat.tech>
Date: Sun, 12 Mar 2023 19:32:20 -0700
Subject: [PATCH] feat(memory): significantly reduce memory use during
 resolution (#203)

Fixes: https://github.com/orogene/orogene/issues/202

With this patch, memory use during resolution is significantly decreased. There's a small perf sacrifice, but I think it's fine and very worth it.

| pm | max res memory |
|---|---|
| oro (w/ this PR) | 274.5 mb |
| oro (pre-patch) | 724.5 mb |
| Yarn | 663.3 mb |
| pnpm | 943.8 mb |
| NPM | 1,028.8 mb |
| bun | 2,761.6 mb |
---
 crates/nassun/src/client.rs    | 11 +++++++++++
 crates/nassun/src/fetch/npm.rs | 26 ++++++++++++++++++++------
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/crates/nassun/src/client.rs b/crates/nassun/src/client.rs
index e847d3c3..beb36f0c 100644
--- a/crates/nassun/src/client.rs
+++ b/crates/nassun/src/client.rs
@@ -27,6 +27,7 @@ pub struct NassunOpts {
     base_dir: Option<PathBuf>,
     default_tag: Option<String>,
     registries: HashMap<Option<String>, Url>,
+    memoize_metadata: bool,
 }
 
 impl NassunOpts {
@@ -65,6 +66,15 @@ impl NassunOpts {
         self
     }
 
+    /// Whether to memoize package metadata. This will keep any processed
+    /// packuments in memory for the lifetime of this `Nassun` instance.
+    /// Setting this to `true` may increase performance when fetching many
+    /// packages, at the cost of significant additional memory usage.
+    pub fn memoize_metadata(mut self, memoize: bool) -> Self {
+        self.memoize_metadata = memoize;
+        self
+    }
+
     /// Build a new Nassun instance from this options object.
     pub fn build(self) -> Nassun {
         let registry = self
@@ -104,6 +114,7 @@ impl NassunOpts {
                 #[allow(clippy::redundant_clone)]
                 client.clone(),
                 self.registries,
+                self.memoize_metadata,
             )),
             #[cfg(not(target_arch = "wasm32"))]
             dir_fetcher: Arc::new(DirFetcher::new()),
diff --git a/crates/nassun/src/fetch/npm.rs b/crates/nassun/src/fetch/npm.rs
index 7d5b1465..feabdfe4 100644
--- a/crates/nassun/src/fetch/npm.rs
+++ b/crates/nassun/src/fetch/npm.rs
@@ -18,17 +18,23 @@ use crate::resolver::PackageResolution;
 pub(crate) struct NpmFetcher {
     client: OroClient,
     registries: HashMap<Option<String>, Url>,
+    cache_packuments: bool,
     packuments: DashMap<String, Arc<Packument>>,
     corgi_packuments: DashMap<String, Arc<CorgiPackument>>,
 }
 
 impl NpmFetcher {
-    pub(crate) fn new(client: OroClient, registries: HashMap<Option<String>, Url>) -> Self {
+    pub(crate) fn new(
+        client: OroClient,
+        registries: HashMap<Option<String>, Url>,
+        cache_packuments: bool,
+    ) -> Self {
         Self {
             client,
             registries,
             packuments: DashMap::new(),
             corgi_packuments: DashMap::new(),
+            cache_packuments,
         }
     }
 }
@@ -99,12 +105,16 @@ impl PackageFetcher for NpmFetcher {
         } = spec.target()
         {
             if let Some(packument) = self.corgi_packuments.get(name) {
-                return Ok(packument.value().clone());
+                if self.cache_packuments {
+                    return Ok(packument.value().clone());
+                }
             }
             let client = self.client.with_registry(self.pick_registry(scope));
             let packument = Arc::new(client.corgi_packument(&name).await?);
-            self.corgi_packuments
-                .insert(name.clone(), packument.clone());
+            if self.cache_packuments {
+                self.corgi_packuments
+                    .insert(name.clone(), packument.clone());
+            }
             Ok(packument)
         } else {
             unreachable!("How did a non-Npm resolution get here?");
@@ -126,11 +136,15 @@ impl PackageFetcher for NpmFetcher {
         } = pkg
         {
             if let Some(packument) = self.packuments.get(name) {
-                return Ok(packument.value().clone());
+                if self.cache_packuments {
+                    return Ok(packument.value().clone());
+                }
             }
             let client = self.client.with_registry(self.pick_registry(scope));
             let packument = Arc::new(client.packument(&name).await?);
-            self.packuments.insert(name.clone(), packument.clone());
+            if self.cache_packuments {
+                self.packuments.insert(name.clone(), packument.clone());
+            }
             Ok(packument)
         } else {
             unreachable!()