feat: Add --log-compute-units flag for compute unit logs

cli/src/lib.rs

@@ -152,6 +152,9 @@ pub enum Command {
         /// Architecture to use when building the program
         #[clap(value_enum, long, default_value = "sbf")]
         arch: ProgramArch,
+        /// Add compute unit logging at key points in the program
+        #[clap(long)]
+        log_compute_units: bool,
     },
     /// Expands macros (wrapper around cargo expand)
     ///
@@ -234,6 +237,9 @@ pub enum Command {
         /// Arguments to pass to the underlying `cargo build-sbf` command.
         #[clap(required = false, last = true)]
         cargo_args: Vec<String>,
+        /// Add compute unit logging at key points in the program
+        #[clap(long)]
+        log_compute_units: bool,
     },
     /// Creates a new program.
     New {
@@ -1117,6 +1123,42 @@ fn get_npm_init_license() -> Result<String> {
     Ok(license.trim().to_string())
 }
 
+fn ensure_program_feature(
+    cfg_override: &ConfigOverride,
+    program_name: Option<String>,
+    feature: &str,
+    flag_name: &str,
+) -> Result<()> {
+    let cfg = Config::discover(cfg_override)?
+        .ok_or_else(|| anyhow!("The `{flag_name}` flag requires an Anchor workspace."))?;
+    let programs = cfg.get_programs(program_name)?;
+    let mut missing = Vec::new();
+
+    for program in programs {
+        let cargo_toml = program.path.join("Cargo.toml");
+        let manifest = Manifest::from_path(&cargo_toml)?;
+        if !manifest.features.contains_key(feature) {
+            missing.push((program.lib_name, cargo_toml));
+        }
+    }
+
+    if missing.is_empty() {
+        return Ok(());
+    }
+
+    let mut message = format!(
+        "`{flag_name}` requires each program to define the `{feature}` feature in its Cargo.toml.\nMissing in:\n"
+    );
+    for (name, path) in missing {
+        message.push_str(&format!("  - {name} ({})\n", path.display()));
+    }
+    message.push_str(&format!(
+        "\nAdd to each program's Cargo.toml:\n\n[features]\n{feature} = []\n\nOr run without `{flag_name}`."
+    ));
+
+    Err(anyhow!(message))
+}
+
 fn process_command(opts: Opts) -> Result<()> {
     match opts.command {
         Command::Init {
@@ -1159,25 +1201,39 @@ fn process_command(opts: Opts) -> Result<()> {
             ignore_keys,
             no_docs,
             arch,
-        } => build(
-            &opts.cfg_override,
-            no_idl,
-            idl,
-            idl_ts,
-            verifiable,
-            skip_lint,
-            ignore_keys,
-            program_name,
-            solana_version,
-            docker_image,
-            bootstrap,
-            None,
-            None,
-            env,
-            cargo_args,
-            no_docs,
-            arch,
-        ),
+            log_compute_units,
+        } => {
+            let mut cargo_args = cargo_args;
+            if log_compute_units {
+                ensure_program_feature(
+                    &opts.cfg_override,
+                    program_name.clone(),
+                    "log-compute-units",
+                    "--log-compute-units",
+                )?;
+                cargo_args.push("--features".to_string());
+                cargo_args.push("log-compute-units".to_string());
+            }
+            build(
+                &opts.cfg_override,
+                no_idl,
+                idl,
+                idl_ts,
+                verifiable,
+                skip_lint,
+                ignore_keys,
+                program_name,
+                solana_version,
+                docker_image,
+                bootstrap,
+                None,
+                None,
+                env,
+                cargo_args,
+                no_docs,
+                arch,
+            )
+        }
         Command::Verify {
             program_id,
             repo_url,
@@ -1252,22 +1308,38 @@ fn process_command(opts: Opts) -> Result<()> {
             cargo_args,
             skip_lint,
             arch,
-        } => test(
-            &opts.cfg_override,
-            program_name,
-            skip_deploy,
-            skip_local_validator,
-            skip_build,
-            skip_lint,
-            no_idl,
-            detach,
-            run,
-            validator,
-            args,
-            env,
-            cargo_args,
-            arch,
-        ),
+            log_compute_units,
+        } => {
+            let mut cargo_args = cargo_args;
+            if log_compute_units {
+                if !skip_build {
+                    ensure_program_feature(
+                        &opts.cfg_override,
+                        program_name.clone(),
+                        "log-compute-units",
+                        "--log-compute-units",
+                    )?;
+                }
+                cargo_args.push("--features".to_string());
+                cargo_args.push("log-compute-units".to_string());
+            }
+            test(
+                &opts.cfg_override,
+                program_name,
+                skip_deploy,
+                skip_local_validator,
+                skip_build,
+                skip_lint,
+                no_idl,
+                detach,
+                run,
+                validator,
+                args,
+                env,
+                cargo_args,
+                arch,
+            )
+        }
         Command::Airdrop { amount, pubkey } => airdrop(&opts.cfg_override, amount, pubkey),
         Command::Cluster { subcmd } => cluster(subcmd),
         Command::Config { subcmd } => config_cmd(&opts.cfg_override, subcmd),

cli/src/rust_template.rs

@@ -217,6 +217,7 @@ cpi = ["no-entrypoint"]
 no-entrypoint = []
 no-idl = []
 no-log-ix-name = []
+log-compute-units = []
 idl-build = ["anchor-lang/idl-build"]
 anchor-debug = []
 custom-heap = []

lang/src/lib.rs

@@ -132,6 +132,25 @@ pub mod solana_program {
             #[cfg(not(target_os = "solana"))]
             core::hint::black_box(data);
         }
+
+        /// Print the remaining compute units available to the program.
+        pub fn sol_log_compute_units() {
+            #[cfg(target_os = "solana")]
+            unsafe {
+                solana_define_syscall::definitions::sol_log_compute_units_()
+            };
+        }
+
+        /// Get the remaining compute units available to the program.
+        pub fn sol_remaining_compute_units() -> u64 {
+            #[cfg(target_os = "solana")]
+            unsafe {
+                solana_define_syscall::definitions::sol_remaining_compute_units()
+            }
+
+            #[cfg(not(target_os = "solana"))]
+            0
+        }
     }
     pub mod sysvar {
         pub use solana_sysvar_id::{declare_deprecated_sysvar_id, declare_sysvar_id, SysvarId};
@@ -502,6 +521,7 @@ pub mod prelude {
     pub use crate as anchor_lang;
     pub use crate::solana_program::account_info::{next_account_info, AccountInfo};
     pub use crate::solana_program::instruction::AccountMeta;
+    pub use crate::solana_program::log::sol_remaining_compute_units;
     pub use crate::solana_program::program_error::ProgramError;
     pub use crate::solana_program::pubkey::Pubkey;
     pub use crate::solana_program::*;

lang/syn/src/codegen/program/entry.rs

@@ -37,6 +37,12 @@ pub fn generate(program: &Program) -> proc_macro2::TokenStream {
         /// The `entry` function here, defines the standard entry to a Solana
         /// program, where execution begins.
         pub fn entry<'info>(program_id: &Pubkey, accounts: &'info [AccountInfo<'info>], data: &[u8]) -> anchor_lang::solana_program::entrypoint::ProgramResult {
+            // Log compute units at program start. Note: This won't be exactly 200,000 because
+            // the Solana runtime consumes CUs for tx deserialization, program loading,
+            // account loading, and signature verification before calling your program.
+            #[cfg(feature = "log-compute-units")]
+            anchor_lang::prelude::msg!("anchor-compute: program-start, {} units", anchor_lang::solana_program::log::sol_remaining_compute_units());
+
             try_entry(program_id, accounts, data).map_err(|e| {
                 e.log();
                 e.into()
@@ -48,6 +54,10 @@ pub fn generate(program: &Program) -> proc_macro2::TokenStream {
             {
                 msg!("anchor-debug is active");
             }
+
+            #[cfg(feature = "log-compute-units")]
+            anchor_lang::prelude::msg!("anchor-compute: entry, {} units", anchor_lang::solana_program::log::sol_remaining_compute_units());
+
             if *program_id != ID {
                 return Err(anchor_lang::error::ErrorCode::DeclaredProgramIdMismatch.into());
             }

lang/syn/src/codegen/program/handlers.rs

@@ -107,12 +107,20 @@ pub fn generate(program: &Program) -> proc_macro2::TokenStream {
                     #[cfg(not(feature = "no-log-ix-name"))]
                     anchor_lang::prelude::msg!(#ix_name_log);
 
+                    // Log compute units at handler start (before any deserialization)
+                    #[cfg(feature = "log-compute-units")]
+                    anchor_lang::prelude::msg!("anchor-compute: begin, {} units", anchor_lang::solana_program::log::sol_remaining_compute_units());
+
                     #param_validation
                     // Deserialize data.
                     let ix = instruction::#ix_name::deserialize(&mut &__ix_data[..])
                         .map_err(|_| anchor_lang::error::ErrorCode::InstructionDidNotDeserialize)?;
                     let instruction::#variant_arm = ix;
 
+                    // Log compute units after instruction data deserialization
+                    #[cfg(feature = "log-compute-units")]
+                    anchor_lang::prelude::msg!("anchor-compute: ix-deser, {} units", anchor_lang::solana_program::log::sol_remaining_compute_units());
+
                     // Bump collector.
                     let mut __bumps = <#anchor as anchor_lang::Bumps>::Bumps::default();
 
@@ -128,6 +136,10 @@ pub fn generate(program: &Program) -> proc_macro2::TokenStream {
                         &mut __reallocs,
                     )?;
 
+                    // Log compute units after account deserialization (often most expensive)
+                    #[cfg(feature = "log-compute-units")]
+                    anchor_lang::prelude::msg!("anchor-compute: accts-deser, {} units", anchor_lang::solana_program::log::sol_remaining_compute_units());
+
                     // Invoke user defined handler.
                     let result = #program_name::#ix_method_name(
                         anchor_lang::context::Context::new(
@@ -139,6 +151,10 @@ pub fn generate(program: &Program) -> proc_macro2::TokenStream {
                         #(#ix_arg_names),*
                     )?;
 
+                    // Log compute units after user handler execution
+                    #[cfg(feature = "log-compute-units")]
+                    anchor_lang::prelude::msg!("anchor-compute: handler, {} units", anchor_lang::solana_program::log::sol_remaining_compute_units());
+
                     // Maybe set Solana return data.
                     #maybe_set_return_data
 

tests/log-compute-units/Anchor.toml

@@ -0,0 +1,20 @@
+[toolchain]
+
+[features]
+resolution = true
+skip-lint = false
+
+[programs.localnet]
+log_compute_units = "848vVGVxp1kcRABrBid3uGCu1iiuPAHFuwpYfMMPPVZ7"
+
+[registry]
+url = "https://api.apr.dev"
+
+[provider]
+cluster = "localnet"
+wallet = "~/.config/solana/id.json"
+
+[scripts]
+test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/**/*.ts"
+
+[hooks]

tests/log-compute-units/Cargo.toml

@@ -0,0 +1,8 @@
+[workspace]
+members = [
+    "programs/*"
+]
+resolver = "2"
+
+[profile.release]
+overflow-checks = true

tests/log-compute-units/package.json

@@ -0,0 +1,19 @@
+{
+  "name": "log-compute-units",
+  "version": "0.32.1",
+  "license": "(MIT OR Apache-2.0)",
+  "homepage": "https://github.com/coral-xyz/anchor#readme",
+  "bugs": {
+    "url": "https://github.com/coral-xyz/anchor/issues"
+  },
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/coral-xyz/anchor.git"
+  },
+  "engines": {
+    "node": ">=17"
+  },
+  "scripts": {
+    "test": "anchor test --skip-lint"
+  }
+}

tests/log-compute-units/programs/log-compute-units/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "log-compute-units"
+version = "0.1.0"
+description = "Test for log-compute-units feature"
+edition = "2021"
+
+[lib]
+crate-type = ["cdylib", "lib"]
+name = "log_compute_units"
+
+[features]
+default = ["log-compute-units"]
+no-entrypoint = []
+cpi = ["no-entrypoint"]
+no-idl = []
+no-log-ix-name = []
+log-compute-units = []
+idl-build = ["anchor-lang/idl-build"]
+anchor-debug = []
+custom-heap = []
+custom-panic = []
+
+[dependencies]
+anchor-lang = { path = "../../../../lang" }