From 67b440a73c845bdcb23bb636bae2576269d5fdfc Mon Sep 17 00:00:00 2001
From: Miles Johnson <milesj@users.noreply.github.com>
Date: Thu, 30 Nov 2023 14:36:28 -0800
Subject: [PATCH] fix: Preload schema plugin for `proto use`. (#312)

---
 CHANGELOG.md                           |  7 +++++++
 crates/cli/src/commands/install_all.rs |  8 ++++++++
 crates/core/src/tool_loader.rs         | 26 ++++++++++++++++++--------
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 314494b13..12f45afb1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,13 @@
 - [Rust](https://github.com/moonrepo/rust-plugin/blob/master/CHANGELOG.md)
 - [TOML schema](https://github.com/moonrepo/schema-plugin/blob/master/CHANGELOG.md)
 
+## Unreleased
+
+#### 🐞 Fixes
+
+- Fixed an issue where `proto use` (or parallel processes) would run into file system
+  collisions when attempting to download and install multiple TOML schema based tools.
+
 ## 0.23.5
 
 #### 🚀 Updates
diff --git a/crates/cli/src/commands/install_all.rs b/crates/cli/src/commands/install_all.rs
index 070087da2..0c3e9748a 100644
--- a/crates/cli/src/commands/install_all.rs
+++ b/crates/cli/src/commands/install_all.rs
@@ -6,6 +6,7 @@ use crate::{
     commands::install::{internal_install, InstallArgs},
 };
 use miette::IntoDiagnostic;
+use proto_core::load_schema_plugin;
 use starbase::system;
 use starbase_styles::color;
 use std::{env, process};
@@ -53,6 +54,13 @@ pub async fn install_all() {
 
     disable_progress_bars();
 
+    // Download the schema plugin before installing tools.
+    // We must do this here, otherwise when multiple schema
+    // based tools are installed in parallel, they will
+    // collide when attempting to download the schema plugin!
+    load_schema_plugin(&loader.proto, &loader.user_config).await?;
+
+    // Then install each tool in parallel!
     let mut futures = vec![];
 
     for tool in tools {
diff --git a/crates/core/src/tool_loader.rs b/crates/core/src/tool_loader.rs
index deb50f719..df234d38e 100644
--- a/crates/core/src/tool_loader.rs
+++ b/crates/core/src/tool_loader.rs
@@ -7,6 +7,7 @@ use extism::{manifest::Wasm, Manifest};
 use miette::IntoDiagnostic;
 use proto_pdk_api::{HostArch, HostEnvironment, HostOS, UserConfigSettings};
 use starbase_utils::{json, toml};
+use std::path::PathBuf;
 use std::{env, path::Path};
 use tracing::{debug, trace};
 use warpgate::{create_http_client_with_options, to_virtual_path, Id, PluginLocator};
@@ -113,6 +114,22 @@ pub fn locate_tool(
     Ok(locator)
 }
 
+pub async fn load_schema_plugin(
+    proto: impl AsRef<ProtoEnvironment>,
+    user_config: &UserConfig,
+) -> miette::Result<PathBuf> {
+    let proto = proto.as_ref();
+    let http_client = create_http_client_with_options(user_config.http.clone())?;
+    let plugin_loader = proto.get_plugin_loader();
+
+    let schema_id = Id::raw(SCHEMA_PLUGIN_KEY);
+    let schema_locator = locate_tool(&schema_id, proto, user_config, true)?;
+
+    plugin_loader
+        .load_plugin_with_client(schema_id, schema_locator, &http_client)
+        .await
+}
+
 pub async fn load_tool_from_locator(
     id: impl AsRef<Id>,
     proto: impl AsRef<ProtoEnvironment>,
@@ -138,16 +155,9 @@ pub async fn load_tool_from_locator(
     {
         debug!(source = ?plugin_path, "Loading TOML plugin");
 
-        let schema_id = Id::raw(SCHEMA_PLUGIN_KEY);
-        let schema_locator = locate_tool(&schema_id, proto, user_config, true)?;
-
         let mut manifest = Tool::create_plugin_manifest(
             proto,
-            Wasm::file(
-                plugin_loader
-                    .load_plugin_with_client(schema_id, schema_locator, &http_client)
-                    .await?,
-            ),
+            Wasm::file(load_schema_plugin(proto, user_config).await?),
         )?;
 
         // Convert TOML to JSON