refa(cargo-shuttle): find binary target based on macro in source (#2059)

* refa(cargo-shuttle): find binary target to build based on macro in source

* add tests

* nits

* nit: clarity

* nit

* fix: don't require final ] when looking for macro

Author: jonaro00

cargo-shuttle/src/builder.rs

@@ -1,71 +1,27 @@
-use std::fs::read_to_string;
+use std::fs;
 use std::path::{Path, PathBuf};
 use std::process::Stdio;
 
-use anyhow::{anyhow, bail, Context};
-use cargo_metadata::{Metadata, Package};
+use anyhow::{bail, Context, Result};
+use cargo_metadata::{Metadata, Package, Target};
 use shuttle_common::constants::RUNTIME_NAME;
 use tokio::io::AsyncBufReadExt;
-use tracing::{debug, error, info, trace};
+use tracing::{error, trace};
 
-#[derive(Clone, Debug, Eq, PartialEq)]
 /// This represents a compiled Shuttle service
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub struct BuiltService {
     pub workspace_path: PathBuf,
-    pub manifest_path: PathBuf,
-    pub package_name: String,
+    pub target_name: String,
     pub executable_path: PathBuf,
 }
 
-impl BuiltService {
-    /// The directory that contains the crate (that Cargo.toml is in)
-    pub fn crate_directory(&self) -> &Path {
-        self.manifest_path
-            .parent()
-            .expect("manifest to be in a directory")
-    }
-
-    /// Try to get the service name of a crate from Shuttle.toml in the crate root, if it doesn't
-    /// exist get it from the Cargo.toml package name of the crate.
-    pub fn service_name(&self) -> anyhow::Result<String> {
-        let shuttle_toml_path = self.crate_directory().join("Shuttle.toml");
-
-        match extract_shuttle_toml_name(shuttle_toml_path) {
-            Ok(service_name) => Ok(service_name),
-            Err(error) => {
-                debug!(?error, "failed to get service name from Shuttle.toml");
-
-                // Couldn't get name from Shuttle.toml, use package name instead.
-                Ok(self.package_name.clone())
-            }
-        }
-    }
-}
-
-fn extract_shuttle_toml_name(path: PathBuf) -> anyhow::Result<String> {
-    let shuttle_toml =
-        read_to_string(path.as_path()).map_err(|_| anyhow!("{} not found", path.display()))?;
-
-    let toml: toml::Value =
-        toml::from_str(&shuttle_toml).context("failed to parse Shuttle.toml")?;
-
-    let name = toml
-        .get("name")
-        .context("couldn't find `name` key in Shuttle.toml")?
-        .as_str()
-        .context("`name` key in Shuttle.toml must be a string")?
-        .to_string();
-
-    Ok(name)
-}
-
-/// Given a project directory path, builds the crate
+/// Builds Shuttle service in given project directory
 pub async fn build_workspace(
     project_path: &Path,
     release_mode: bool,
     tx: tokio::sync::mpsc::Sender<String>,
-    deployment: bool,
-) -> anyhow::Result<Vec<BuiltService>> {
+) -> Result<BuiltService> {
     let project_path = project_path.to_owned();
     let manifest_path = project_path.join("Cargo.toml");
     if !manifest_path.exists() {
@@ -100,29 +56,23 @@ pub async fn build_workspace(
     notification.abort();
 
     let metadata = async_cargo_metadata(manifest_path.as_path()).await?;
-    let packages = find_shuttle_packages(&metadata)?;
-    if packages.is_empty() {
-        bail!(
-            "Did not find any packages that Shuttle can run. \
-            Make sure your crate has a binary target that uses `#[shuttle_runtime::main]`."
-        );
-    }
+    let (package, target) = find_first_shuttle_package(&metadata)?;
 
-    let services = compile(
-        packages,
+    let service = cargo_build(
+        package,
+        target,
         release_mode,
         project_path.clone(),
         metadata.target_directory.clone(),
-        deployment,
         tx.clone(),
     )
     .await?;
-    trace!("packages compiled");
+    trace!("package compiled");
 
-    Ok(services)
+    Ok(service)
 }
 
-pub async fn async_cargo_metadata(manifest_path: &Path) -> anyhow::Result<Metadata> {
+pub async fn async_cargo_metadata(manifest_path: &Path) -> Result<Metadata> {
     let metadata = {
         // Modified implementaion of `cargo_metadata::MetadataCommand::exec` (from v0.15.3).
         // Uses tokio Command instead of std, to make this operation non-blocking.
@@ -149,63 +99,62 @@ pub async fn async_cargo_metadata(manifest_path: &Path) -> anyhow::Result<Metada
     Ok(metadata)
 }
 
-pub fn find_shuttle_packages(metadata: &Metadata) -> anyhow::Result<Vec<Package>> {
+/// Find crates with a runtime dependency and main macro
+fn find_shuttle_packages(metadata: &Metadata) -> Result<Vec<(Package, Target)>> {
     let mut packages = Vec::new();
+    trace!("Finding Shuttle-related packages");
     for member in metadata.workspace_packages() {
-        // skip non-Shuttle-related crates
-        // (must have runtime dependency and not be just a library)
         let has_runtime_dep = member
             .dependencies
             .iter()
             .any(|dependency| dependency.name == RUNTIME_NAME);
-        let has_binary_target = member.targets.iter().any(|t| t.is_bin());
-        if !(has_runtime_dep && has_binary_target) {
+        if !has_runtime_dep {
+            trace!("Skipping {}, no shuttle-runtime dependency", member.name);
             continue;
         }
 
-        let mut shuttle_deps = member
-            .dependencies
-            .iter()
-            .filter(|&d| d.name.starts_with("shuttle-"))
-            .map(|d| format!("{} '{}'", d.name, d.req))
-            .collect::<Vec<_>>();
-        shuttle_deps.sort();
-        info!(name = member.name, deps = ?shuttle_deps, "Found workspace member with shuttle dependencies");
-        packages.push(member.to_owned());
+        let mut target = None;
+        for t in member.targets.iter() {
+            if t.is_bin()
+                && fs::read_to_string(t.src_path.as_std_path())
+                    .context("reading to check for shuttle macro")?
+                    // don't look for the last ] so that we also find instances of `#[shuttle_runtime::main(...)]`
+                    .contains("#[shuttle_runtime::main")
+            {
+                target = Some(t);
+                break;
+            }
+        }
+        let Some(target) = target else {
+            trace!(
+                "Skipping {}, no binary target with a #[shuttle_runtime::main] macro",
+                member.name
+            );
+            continue;
+        };
+
+        trace!("Found {}", member.name);
+        packages.push((member.to_owned(), target.to_owned()));
     }
 
     Ok(packages)
 }
 
-// Only used in deployer
-pub async fn clean_crate(project_path: &Path) -> anyhow::Result<()> {
-    let manifest_path = project_path.join("Cargo.toml");
-    if !manifest_path.exists() {
-        bail!("Cargo manifest file not found: {}", manifest_path.display());
-    }
-    if !tokio::process::Command::new("cargo")
-        .arg("clean")
-        .arg("--manifest-path")
-        .arg(&manifest_path)
-        .arg("--offline")
-        .status()
-        .await?
-        .success()
-    {
-        bail!("`cargo clean` failed. Did you build anything yet?");
-    }
-
-    Ok(())
+pub fn find_first_shuttle_package(metadata: &Metadata) -> Result<(Package, Target)> {
+    find_shuttle_packages(metadata)?.into_iter().next().context(
+        "Expected at least one target that Shuttle can build. \
+        Make sure your crate has a binary target that uses a fully qualified `#[shuttle_runtime::main]`.",
+    )
 }
 
-async fn compile(
-    packages: Vec<Package>,
+async fn cargo_build(
+    package: Package,
+    target: Target,
     release_mode: bool,
     project_path: PathBuf,
     target_path: impl Into<PathBuf>,
-    deployment: bool,
     tx: tokio::sync::mpsc::Sender<String>,
-) -> anyhow::Result<Vec<BuiltService>> {
+) -> Result<BuiltService> {
     let manifest_path = project_path.join("Cargo.toml");
     if !manifest_path.exists() {
         bail!("Cargo manifest file not found: {}", manifest_path.display());
@@ -221,17 +170,11 @@ async fn compile(
         .arg("--color=always") // piping disables auto color, but we want it
         .current_dir(project_path.as_path());
 
-    if deployment {
-        cmd.arg("--jobs=4");
-    }
-
-    // TODO: Compile only one binary target in the package.
-    for package in &packages {
-        if package.features.contains_key("shuttle") {
-            cmd.arg("--no-default-features").arg("--features=shuttle");
-        }
-        cmd.arg("--package").arg(package.name.as_str());
+    if package.features.contains_key("shuttle") {
+        cmd.arg("--no-default-features").arg("--features=shuttle");
     }
+    cmd.arg("--package").arg(package.name.as_str());
+    cmd.arg("--bin").arg(target.name.as_str());
 
     let profile = if release_mode {
         cmd.arg("--release");
@@ -255,30 +198,22 @@ async fn compile(
     });
     let status = handle.wait().await?;
     if !status.success() {
-        bail!("Build failed. Is the Shuttle runtime missing?");
+        bail!("Build failed.");
     }
 
-    let services = packages
-        .iter()
-        .map(|package| {
-            let mut path: PathBuf = [
-                project_path.clone(),
-                target_path.clone(),
-                profile.into(),
-                package.name.clone().into(),
-            ]
-            .iter()
-            .collect();
-            path.set_extension(std::env::consts::EXE_EXTENSION);
-
-            BuiltService {
-                workspace_path: project_path.clone(),
-                manifest_path: package.manifest_path.clone().into_std_path_buf(),
-                package_name: package.name.clone(),
-                executable_path: path,
-            }
-        })
-        .collect();
-
-    Ok(services)
+    let mut path: PathBuf = [
+        project_path.clone(),
+        target_path.clone(),
+        profile.into(),
+        target.name.as_str().into(),
+    ]
+    .iter()
+    .collect();
+    path.set_extension(std::env::consts::EXE_EXTENSION);
+
+    Ok(BuiltService {
+        workspace_path: project_path.clone(),
+        target_name: target.name,
+        executable_path: path,
+    })
 }

cargo-shuttle/src/lib.rs

@@ -61,7 +61,9 @@ use crate::args::{
     ResourceCommand, SecretsArgs, TableArgs, TemplateLocation,
 };
 pub use crate::args::{Command, ProjectArgs, RunArgs, ShuttleArgs};
-use crate::builder::{async_cargo_metadata, build_workspace, find_shuttle_packages, BuiltService};
+use crate::builder::{
+    async_cargo_metadata, build_workspace, find_first_shuttle_package, BuiltService,
+};
 use crate::config::RequestContext;
 use crate::provisioner_server::{ProvApiState, ProvisionerServer};
 use crate::util::{
@@ -1293,7 +1295,7 @@ impl Shuttle {
         })
     }
 
-    async fn pre_local_run(&self, run_args: &RunArgs) -> Result<Vec<BuiltService>> {
+    async fn pre_local_run(&self, run_args: &RunArgs) -> Result<BuiltService> {
         trace!("starting a local run with args: {run_args:?}");
 
         let (tx, mut rx) = tokio::sync::mpsc::channel::<String>(256);
@@ -1311,7 +1313,7 @@ impl Shuttle {
             working_directory.display()
         );
 
-        build_workspace(working_directory, run_args.release, tx, false).await
+        build_workspace(working_directory, run_args.release, tx).await
     }
 
     fn find_available_port(run_args: &mut RunArgs) {
@@ -1346,12 +1348,7 @@ impl Shuttle {
             return bacon::run_bacon(working_directory).await;
         }
 
-        let services = self.pre_local_run(&run_args).await?;
-        let service = services
-            .first()
-            .expect("at least one shuttle service")
-            .to_owned();
-
+        let service = self.pre_local_run(&run_args).await?;
         trace!(path = ?service.executable_path, "runtime executable");
 
         let secrets = Shuttle::get_secrets(&run_args.secret_args, working_directory, true)?
@@ -1382,7 +1379,7 @@ impl Shuttle {
         println!(
             "\n    {} {} on http://{}:{}\n",
             "Starting".bold().green(),
-            service.package_name,
+            service.target_name,
             ip,
             run_args.port,
         );
@@ -1592,13 +1589,10 @@ impl Shuttle {
         let mut rust_build_args = BuildArgsRust::default();
 
         let metadata = async_cargo_metadata(manifest_path.as_path()).await?;
-        let packages = find_shuttle_packages(&metadata)?;
         // TODO: support overriding this
-        let package = packages
-            .first()
-            .expect("Expected at least one crate with shuttle-runtime in the workspace");
-        let package_name = package.name.to_owned();
-        rust_build_args.package_name = Some(package_name);
+        let (package, target) = find_first_shuttle_package(&metadata)?;
+        rust_build_args.package_name = Some(package.name.clone());
+        rust_build_args.binary_name = Some(target.name.clone());
 
         // activate shuttle feature if present
         let (no_default_features, features) = if package.features.contains_key("shuttle") {