sdk, programs/bpf_loader: add sol_remaining_compute_units syscall (#31640)

bpf_loader: add sol_remaining_compute_units syscall

Co-authored-by: jonch <9093549+jon-chuang@users.noreply.github.com>

docs/src/developing/on-chain-programs/developing-c.md

@@ -157,6 +157,9 @@ with program logs.
 
 ## Compute Budget
 
+Use the system call `sol_remaining_compute_units()` to return a `u64` indicating
+the number of compute units remaining for this transaction.
+
 Use the system call
 [`sol_log_compute_units()`](https://github.com/solana-labs/solana/blob/d3a3a7548c857f26ec2cb10e270da72d373020ec/sdk/sbf/c/inc/solana_sdk.h#L140)
 to log a message containing the remaining number of compute units the program

docs/src/developing/on-chain-programs/developing-rust.md

@@ -355,6 +355,9 @@ fn custom_panic(info: &core::panic::PanicInfo<'_>) {
 
 ## Compute Budget
 
+Use the system call `sol_remaining_compute_units()` to return a `u64` indicating
+the number of compute units remaining for this transaction.
+
 Use the system call
 [`sol_log_compute_units()`](https://github.com/solana-labs/solana/blob/d9b0fc0e3eec67dfe4a97d9298b15969b2804fab/sdk/program/src/log.rs#L141)
 to log a message containing the remaining number of compute units the program

program-runtime/src/compute_budget.rs

@@ -130,6 +130,8 @@ pub struct ComputeBudget {
     /// of compute units consumed to call poseidon syscall for a given number
     /// of inputs.
     pub poseidon_cost_coefficient_c: u64,
+    /// Number of compute units consumed for accessing the remaining compute units.
+    pub get_remaining_compute_units_cost: u64,
 }
 
 impl Default for ComputeBudget {
@@ -181,6 +183,7 @@ impl ComputeBudget {
             loaded_accounts_data_size_limit: MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES,
             poseidon_cost_coefficient_a: 61,
             poseidon_cost_coefficient_c: 542,
+            get_remaining_compute_units_cost: 100,
         }
     }
 

programs/bpf_loader/src/syscalls/mod.rs

@@ -40,8 +40,8 @@ use {
             enable_early_verification_of_account_modifications, enable_partitioned_epoch_reward,
             enable_poseidon_syscall, error_on_syscall_bpf_function_hash_collisions,
             last_restart_slot_sysvar, libsecp256k1_0_5_upgrade_enabled, reject_callx_r10,
-            stop_sibling_instruction_search_at_parent, stop_truncating_strings_in_syscalls,
-            switch_to_new_elf_parser,
+            remaining_compute_units_syscall_enabled, stop_sibling_instruction_search_at_parent,
+            stop_truncating_strings_in_syscalls, switch_to_new_elf_parser,
         },
         hash::{Hasher, HASH_BYTES},
         instruction::{
@@ -164,6 +164,8 @@ pub fn create_program_runtime_environment_v1<'a>(
         && feature_set.is_active(&disable_deploy_of_alloc_free_syscall::id());
     let last_restart_slot_syscall_enabled = feature_set.is_active(&last_restart_slot_sysvar::id());
     let enable_poseidon_syscall = feature_set.is_active(&enable_poseidon_syscall::id());
+    let remaining_compute_units_syscall_enabled =
+        feature_set.is_active(&remaining_compute_units_syscall_enabled::id());
     // !!! ATTENTION !!!
     // When adding new features for RBPF here,
     // also add them to `Bank::apply_builtin_program_feature_transitions()`.
@@ -335,6 +337,14 @@ pub fn create_program_runtime_environment_v1<'a>(
         SyscallPoseidon::call,
     )?;
 
+    // Accessing remaining compute units
+    register_feature_gated_function!(
+        result,
+        remaining_compute_units_syscall_enabled,
+        *b"sol_remaining_compute_units",
+        SyscallRemainingComputeUnits::call
+    )?;
+
     // Log data
     result.register_function_hashed(*b"sol_log_data", SyscallLogData::call)?;
 
@@ -1877,6 +1887,26 @@ declare_syscall!(
     }
 );
 
+declare_syscall!(
+    /// Read remaining compute units
+    SyscallRemainingComputeUnits,
+    fn inner_call(
+        invoke_context: &mut InvokeContext,
+        _arg1: u64,
+        _arg2: u64,
+        _arg3: u64,
+        _arg4: u64,
+        _arg5: u64,
+        _memory_mapping: &mut MemoryMapping,
+    ) -> Result<u64, Error> {
+        let budget = invoke_context.get_compute_budget();
+        consume_compute_meter(invoke_context, budget.syscall_base_cost)?;
+
+        use solana_rbpf::vm::ContextObject;
+        Ok(invoke_context.get_remaining())
+    }
+);
+
 #[cfg(test)]
 #[allow(clippy::arithmetic_side_effects)]
 #[allow(clippy::indexing_slicing)]

programs/sbf/Cargo.lock

@@ -5889,6 +5889,16 @@ dependencies = [
  "solana-sbf-rust-realloc",
 ]
 
+[[package]]
+name = "solana-sbf-rust-remaining-compute-units"
+version = "1.17.0"
+dependencies = [
+ "solana-program",
+ "solana-program-runtime",
+ "solana-program-test",
+ "solana-sdk",
+]
+
 [[package]]
 name = "solana-sbf-rust-ro-account_modify"
 version = "1.17.0"

programs/sbf/Cargo.toml

@@ -145,6 +145,7 @@ members = [
     "rust/rand",
     "rust/realloc",
     "rust/realloc_invoke",
+    "rust/remaining_compute_units",
     "rust/ro_account_modify",
     "rust/ro_modify",
     "rust/sanity",

programs/sbf/build.rs

@@ -96,6 +96,7 @@ fn main() {
             "rand",
             "realloc",
             "realloc_invoke",
+            "remaining_compute_units",
             "ro_modify",
             "ro_account_modify",
             "sanity",

programs/sbf/c/src/remaining_compute_units/remaining_compute_units.c

@@ -0,0 +1,23 @@
+/**
+ * @brief sol_remaining_compute_units Syscall test
+ */
+#include <solana_sdk.h>
+#include <stdio.h>
+
+extern uint64_t entrypoint(const uint8_t *input) {
+  char buffer[200];
+
+  int i = 0;
+  for (; i < 100000; ++i) {
+    if (i % 500 == 0) {
+      uint64_t remaining = sol_remaining_compute_units();
+      snprintf(buffer, 200, "remaining compute units: %d", (int)remaining);
+      sol_log(buffer);
+      if (remaining < 25000) {
+        break;
+      }
+    }
+  }
+
+  return SUCCESS;
+}

programs/sbf/rust/remaining_compute_units/Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "solana-sbf-rust-remaining-compute-units"
+documentation = "https://docs.rs/solana-sbf-rust-remaining-compute-units"
+version = { workspace = true }
+description = { workspace = true }
+authors = { workspace = true }
+repository = { workspace = true }
+homepage = { workspace = true }
+license = { workspace = true }
+edition = { workspace = true }
+
+[features]
+no-entrypoint = []
+test-bpf = []
+dummy-for-ci-check = ["test-bpf"]
+
+[dependencies]
+solana-program = { workspace = true }
+
+[dev-dependencies]
+solana-program-runtime = { workspace = true }
+solana-program-test = { workspace = true }
+solana-sdk = { workspace = true }
+
+[lib]
+crate-type = ["cdylib", "lib"]
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]

programs/sbf/rust/remaining_compute_units/src/lib.rs

@@ -0,0 +1,29 @@
+//! @brief Example Rust-based BPF program that exercises the sol_remaining_compute_units syscall
+
+extern crate solana_program;
+use solana_program::{
+    account_info::AccountInfo, compute_units::sol_remaining_compute_units,
+    entrypoint::ProgramResult, msg, pubkey::Pubkey,
+};
+solana_program::entrypoint!(process_instruction);
+pub fn process_instruction(
+    _program_id: &Pubkey,
+    _accounts: &[AccountInfo],
+    _instruction_data: &[u8],
+) -> ProgramResult {
+    let mut i = 0u32;
+    for _ in 0..100_000 {
+        if i % 500 == 0 {
+            let remaining = sol_remaining_compute_units();
+            msg!("remaining compute units: {:?}", remaining);
+            if remaining < 25_000 {
+                break;
+            }
+        }
+        i = i.saturating_add(1);
+    }
+
+    msg!("i: {:?}", i);
+
+    Ok(())
+}

programs/sbf/rust/remaining_compute_units/tests/lib.rs

@@ -0,0 +1,27 @@
+#![cfg(feature = "test-bpf")]
+
+use {
+    solana_program_test::*,
+    solana_sbf_rust_remaining_compute_units::process_instruction,
+    solana_sdk::{
+        instruction::Instruction, pubkey::Pubkey, signature::Signer, transaction::Transaction,
+    },
+};
+
+#[tokio::test]
+async fn test_remaining_compute_units() {
+    let program_id = Pubkey::new_unique();
+    let program_test = ProgramTest::new(
+        "solana_sbf_rust_remaining_compute_units",
+        program_id,
+        processor!(process_instruction),
+    );
+    let (mut banks_client, payer, recent_blockhash) = program_test.start().await;
+
+    let mut transaction = Transaction::new_with_payer(
+        &[Instruction::new_with_bincode(program_id, &(), vec![])],
+        Some(&payer.pubkey()),
+    );
+    transaction.sign(&[&payer], recent_blockhash);
+    banks_client.process_transaction(transaction).await.unwrap();
+}

runtime/src/bank.rs

@@ -8111,6 +8111,7 @@ impl Bank {
             feature_set::disable_deploy_of_alloc_free_syscall::id(),
             feature_set::last_restart_slot_sysvar::id(),
             feature_set::delay_visibility_of_program_deployment::id(),
+            feature_set::remaining_compute_units_syscall_enabled::id(),
         ];
         if !only_apply_transitions_for_new_features
             || FEATURES_AFFECTING_RBPF

sdk/program/src/compute_units.rs

@@ -0,0 +1,13 @@
+/// Return the remaining compute units the program may consume
+#[inline]
+pub fn sol_remaining_compute_units() -> u64 {
+    #[cfg(target_os = "solana")]
+    unsafe {
+        crate::syscalls::sol_remaining_compute_units()
+    }
+
+    #[cfg(not(target_os = "solana"))]
+    {
+        crate::program_stubs::sol_remaining_compute_units()
+    }
+}

sdk/program/src/lib.rs

@@ -483,6 +483,7 @@ pub mod bpf_loader;
 pub mod bpf_loader_deprecated;
 pub mod bpf_loader_upgradeable;
 pub mod clock;
+pub mod compute_units;
 pub mod debug_account_data;
 pub mod decode_error;
 pub mod ed25519_program;

sdk/program/src/program_stubs.rs

@@ -30,6 +30,10 @@ pub trait SyscallStubs: Sync + Send {
     fn sol_log_compute_units(&self) {
         sol_log("SyscallStubs: sol_log_compute_units() not available");
     }
+    fn sol_remaining_compute_units(&self) -> u64 {
+        sol_log("SyscallStubs: sol_remaining_compute_units() defaulting to 0");
+        0
+    }
     fn sol_invoke_signed(
         &self,
         _instruction: &Instruction,
@@ -126,6 +130,10 @@ pub(crate) fn sol_log_compute_units() {
     SYSCALL_STUBS.read().unwrap().sol_log_compute_units();
 }
 
+pub(crate) fn sol_remaining_compute_units() -> u64 {
+    SYSCALL_STUBS.read().unwrap().sol_remaining_compute_units()
+}
+
 pub(crate) fn sol_invoke_signed(
     instruction: &Instruction,
     account_infos: &[AccountInfo],

sdk/program/src/syscalls/definitions.rs

@@ -70,6 +70,7 @@ define_syscall!(fn sol_alt_bn128_group_op(group_op: u64, input: *const u8, input
 define_syscall!(fn sol_big_mod_exp(params: *const u8, result: *mut u8) -> u64);
 define_syscall!(fn sol_get_epoch_rewards_sysvar(addr: *mut u8) -> u64);
 define_syscall!(fn sol_poseidon(parameters: u64, endianness: u64, vals: *const u8, val_len: u64, hash_result: *mut u8) -> u64);
+define_syscall!(fn sol_remaining_compute_units() -> u64);
 
 #[cfg(target_feature = "static-syscalls")]
 pub const fn sys_hash(name: &str) -> usize {

sdk/sbf/c/inc/sol/compute_units.h

@@ -0,0 +1,42 @@
+#pragma once
+/**
+ * @brief Solana logging utilities
+ */
+
+#include <sol/types.h>
+#include <sol/string.h>
+#include <sol/entrypoint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * Prints a string to stdout
+ */
+/* DO NOT MODIFY THIS GENERATED FILE. INSTEAD CHANGE sdk/sbf/c/inc/sol/inc/compute_units.inc AND RUN `cargo run --bin gen-headers` */
+#ifndef SOL_SBFV2
+uint64_t sol_remaining_compute_units();
+#else
+typedef uint64_t(*sol_remaining_compute_units_pointer_type)();
+static uint64_t sol_remaining_compute_units() {
+  sol_remaining_compute_units_pointer_type sol_remaining_compute_units_pointer = (sol_remaining_compute_units_pointer_type) 3991886574;
+  return sol_remaining_compute_units_pointer();
+}
+#endif
+
+#ifdef SOL_TEST
+/**
+ * Stub functions when building tests
+ */
+
+uint64_t sol_remaining_compute_units() {
+  return UINT64_MAX;
+}
+#endif
+
+#ifdef __cplusplus
+}
+#endif
+
+/**@}*/

sdk/sbf/c/inc/sol/inc/compute_units.inc

@@ -0,0 +1,33 @@
+#pragma once
+/**
+ * @brief Solana logging utilities
+ */
+
+#include <sol/types.h>
+#include <sol/string.h>
+#include <sol/entrypoint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * Prints a string to stdout
+ */
+@SYSCALL uint64_t sol_remaining_compute_units();
+
+#ifdef SOL_TEST
+/**
+ * Stub functions when building tests
+ */
+
+uint64_t sol_remaining_compute_units() {
+  return UINT64_MAX;
+}
+#endif
+
+#ifdef __cplusplus
+}
+#endif
+
+/**@}*/

sdk/sbf/c/inc/solana_sdk.h

@@ -6,6 +6,7 @@
 #include <sol/assert.h>
 #include <sol/big_mod_exp.h>
 #include <sol/blake3.h>
+#include <sol/compute_units.h>
 #include <sol/cpi.h>
 #include <sol/deserialize.h>
 #include <sol/deserialize_deprecated.h>

sdk/src/feature_set.rs

@@ -691,6 +691,10 @@ pub mod timely_vote_credits {
     solana_sdk::declare_id!("2oXpeh141pPZCTCFHBsvCwG2BtaHZZAtrVhwaxSy6brS");
 }
 
+pub mod remaining_compute_units_syscall_enabled {
+    solana_sdk::declare_id!("5TuppMutoyzhUSfuYdhgzD47F92GL1g89KpCZQKqedxP");
+}
+
 lazy_static! {
     /// Map of feature identifiers to user-visible description
     pub static ref FEATURE_NAMES: HashMap<Pubkey, &'static str> = [
@@ -856,6 +860,7 @@ lazy_static! {
         (revise_turbine_epoch_stakes::id(), "revise turbine epoch stakes"),
         (enable_poseidon_syscall::id(), "Enable Poseidon syscall"),
         (timely_vote_credits::id(), "use timeliness of votes in determining credits to award"),
+        (remaining_compute_units_syscall_enabled::id(), "enable the remaining_compute_units syscall"),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()