Stop conflating EnergyQuanta and FunctionBudget (#4930)

# Description of Changes

A sort of followup to/unrevert of #4884, an alternative to #4927.

In #3832, we made FunctionBudget a different unit than EnergyQuanta, but
there were still many places where we treated them the same and directly
converted between them. I got confused by that in #4884, and now this PR
properly separates them and corrects the v8 energy calculation.

Also, since we now benchmark the same module in rust and typescript in
the form of the keynote benchmark, this patch adds a new assertion to
that test that the fuel usage of both modules is within 2X of each
other.

# Expected complexity level and risk

3 - touches energy calculation without reverting the problematic #4884,
but I'm confident it's correct this time.

# Testing

- [x] Verified that the conversion factors make sense for wasmtime and
for v8.
- [x] Assert that typescript and rust keynote bench runs have similar
cpu usage

---------

Co-authored-by: joshua-spacetime <josh@clockworklabs.io>

Author: coolreader18

Cargo.lock

@@ -17,12 +17,12 @@ impl EnergyQuanta {
     pub const ZERO: Self = EnergyQuanta { quanta: 0 };
 
     #[inline]
-    pub fn new(quanta: u128) -> Self {
+    pub const fn new(quanta: u128) -> Self {
         Self { quanta }
     }
 
     #[inline]
-    pub fn get(&self) -> u128 {
+    pub const fn get(&self) -> u128 {
         self.quanta
     }
 
@@ -121,46 +121,52 @@ impl fmt::Debug for EnergyBalance {
     }
 }
 
-/// A measure of energy representing the energy budget for a reducer or any callable function.
+/// A measure of the energy budget for a reducer or any callable function.
 ///
-/// In contrast to [`EnergyQuanta`], this is represented by a 64-bit integer. This makes energy handling
-/// for reducers easier, while still providing a unlikely-to-ever-be-reached maximum value (e.g. for wasmtime:
-/// `(u64::MAX eV / 1000 eV/instruction) * 3 ns/instruction = 640 days`)
-#[derive(Copy, Clone, From, Add, Sub, AddAssign, SubAssign)]
+/// This unit is not directly convertible to `EnergyQuanta`. It is currently
+/// 1:1 to wasmtime fuel, and we intend to treat it as representing a CPU
+/// instruction.
+#[derive(Copy, Clone, From, Add, Sub, AddAssign, SubAssign, Debug, PartialEq, Eq, PartialOrd, Ord)]
 pub struct FunctionBudget(u64);
 
 impl FunctionBudget {
-    // 1 second of wasm runtime is roughly 2 TeV, so this is
-    // roughly 1 minute of wasm runtime
-    pub const DEFAULT_BUDGET: Self = FunctionBudget(120_000_000_000_000);
+    /// We've generally assumed that 1 second of wasm runtime uses 2_000_000_000 fuel.
+    /// Currently, 1 wasmtime fuel unit is equivalent to 1 wasm instructions. Assuming
+    /// 1 wasm instruction compiles to 1 CPU instruction (which it doesn't), this implies
+    /// a 1 instruction-per-cycle abstract machine with a CPU frequency of 2GHz.
+    pub const PER_EXECUTION_SEC: Self = FunctionBudget(2_000_000_000);
+
+    pub const PER_EXECUTION_NANOSEC: Self = Self(Self::PER_EXECUTION_SEC.0 / 1_000_000_000);
+
+    /// Roughly 1 minute of runtime.
+    pub const DEFAULT_BUDGET: Self = FunctionBudget(Self::PER_EXECUTION_SEC.0 * 60);
 
     pub const ZERO: Self = FunctionBudget(0);
     pub const MAX: Self = FunctionBudget(u64::MAX);
 
-    pub fn new(v: u64) -> Self {
+    pub const fn new(v: u64) -> Self {
         Self(v)
     }
 
-    pub fn get(&self) -> u64 {
+    pub const fn get(&self) -> u64 {
         self.0
     }
 
-    /// Convert from [`EnergyQuanta`]. Returns `None` if `energy` is too large to be represented.
-    pub fn from_energy(energy: EnergyQuanta) -> Option<Self> {
-        energy.get().try_into().ok().map(Self)
-    }
-}
-
-impl From<FunctionBudget> for EnergyQuanta {
-    fn from(value: FunctionBudget) -> Self {
-        EnergyQuanta::new(value.0.into())
+    /// Convert `FunctionBudget` to `Duration` using the conversion factor
+    /// [`FunctionBudget::PER_EXECUTION_SEC`].
+    pub fn to_duration(self) -> Duration {
+        Duration::from_nanos(self.0 / Self::PER_EXECUTION_NANOSEC.0)
     }
-}
 
-impl fmt::Debug for FunctionBudget {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_tuple("ReducerBudget")
-            .field(&EnergyQuanta::from(*self))
-            .finish()
+    /// Convert `Duration` to `FunctionBudget` using the conversion factor
+    /// [`FunctionBudget::PER_EXECUTION_SEC`].
+    ///
+    /// Returns `None` on overflow, which means that dur > 106.75 days.
+    // in order for duration_nanos * budget_per_ns >= u64::MAX:
+    //              duration_nanos >= u64::MAX / budget_per_ns
+    //              duration_nanos >= (9223372036854775 ns = 106.75 days)
+    pub fn from_duration(dur: Duration) -> Option<Self> {
+        let duration_nanos = u64::try_from(dur.as_nanos()).ok()?;
+        duration_nanos.checked_mul(Self::PER_EXECUTION_NANOSEC.get()).map(Self)
     }
 }

crates/client-api-messages/src/energy.rs

@@ -12,7 +12,7 @@ use spacetimedb::auth::token_validation::{
     new_validator, DefaultValidator, TokenSigner, TokenValidationError, TokenValidator,
 };
 use spacetimedb::auth::JwtKeys;
-use spacetimedb::energy::EnergyQuanta;
+use spacetimedb::energy::FunctionBudget;
 use spacetimedb::identity::Identity;
 use spacetimedb_data_structures::map::HashMap;
 use std::time::{Duration, SystemTime};
@@ -518,7 +518,7 @@ impl headers::Header for SpacetimeIdentityToken {
     }
 }
 
-pub struct SpacetimeEnergyUsed(pub EnergyQuanta);
+pub struct SpacetimeEnergyUsed(pub FunctionBudget);
 impl headers::Header for SpacetimeEnergyUsed {
     fn name() -> &'static http::HeaderName {
         static NAME: http::HeaderName = http::HeaderName::from_static("spacetime-energy-used");
@@ -530,9 +530,7 @@ impl headers::Header for SpacetimeEnergyUsed {
     }
 
     fn encode<E: Extend<HeaderValue>>(&self, values: &mut E) {
-        let mut buf = itoa::Buffer::new();
-        let value = buf.format(self.0.get());
-        values.extend([value.try_into().unwrap()]);
+        values.extend([self.0.get().into()]);
     }
 }
 

crates/client-api/src/auth.rs

@@ -202,7 +202,7 @@ pub async fn call<S: ControlStateDelegate + NodeDelegate>(
             let (status, body) = reducer_outcome_response(&owner_identity, &reducer, result.outcome);
             Ok((
                 status,
-                TypedHeader(SpacetimeEnergyUsed(result.energy_used)),
+                TypedHeader(SpacetimeEnergyUsed(result.execution_budget_used)),
                 TypedHeader(SpacetimeExecutionDurationMicros(result.execution_duration)),
                 body,
             )

crates/client-api/src/routes/database.rs

@@ -1,6 +1,6 @@
 use super::messages::{SubscriptionUpdateMessage, SwitchedServerMessage, ToProtocol, TransactionUpdateMessage};
 use super::{ClientConnection, DataMessage, MessageHandleError, Protocol};
-use crate::energy::EnergyQuanta;
+use crate::energy::FunctionBudget;
 use crate::host::module_host::{EventStatus, ModuleEvent, ModuleFunctionCall};
 use crate::host::{FunctionArgs, ReducerId};
 use crate::identity::Identity;
@@ -139,7 +139,7 @@ impl MessageExecutionError {
             },
             status: EventStatus::FailedInternal(format!("{:#}", err)),
             reducer_return_value: None,
-            energy_quanta_used: EnergyQuanta::ZERO,
+            execution_budget_used: FunctionBudget::ZERO,
             host_execution_duration: Duration::ZERO,
             request_id: Some(RequestId::default()),
             timer: None,

crates/core/src/client/message_handlers_v1.rs

@@ -7,6 +7,7 @@ use crate::subscription::websocket_building::{brotli_compress, decide_compressio
 use bytes::{BufMut, Bytes, BytesMut};
 use bytestring::ByteString;
 use derive_more::From;
+use spacetimedb_client_api_messages::energy::EnergyQuanta;
 use spacetimedb_client_api_messages::websocket::common::{self as ws_common, RowListLen as _};
 use spacetimedb_client_api_messages::websocket::v1::{self as ws_v1};
 use spacetimedb_client_api_messages::websocket::v2 as ws_v2;
@@ -434,7 +435,12 @@ impl ToProtocol for TransactionUpdateMessage {
                     args,
                     request_id,
                 },
-                energy_quanta_used: event.energy_quanta_used,
+                // This conversion is lying. We used to tell the client how much eV a transaction
+                // used, but now the database just tracks cpu usage, and it's converted to energy
+                // elsewhere. So, we just pretend that this is `EnergyQuanta` when it's actually
+                // a different unit, and it doesn't really matter to the client anyway.
+                // TODO(noa): maybe we could just have this be zero, unconditionally?
+                energy_quanta_used: EnergyQuanta::new(event.execution_budget_used.get().into()),
                 total_host_execution_duration: event.host_execution_duration.into(),
                 caller_connection_id: event.caller_connection_id.unwrap_or(ConnectionId::ZERO),
             };

crates/core/src/client/messages.rs

@@ -17,7 +17,7 @@ pub trait EnergyMonitor: Send + Sync + 'static {
     fn record_reducer(
         &self,
         fingerprint: &FunctionFingerprint<'_>,
-        energy_used: EnergyQuanta,
+        execution_budget_used: FunctionBudget,
         execution_duration: Duration,
     );
     fn record_disk_usage(&self, database: &Database, replica_id: u64, disk_usage: u64, period: Duration);
@@ -36,7 +36,7 @@ impl EnergyMonitor for NullEnergyMonitor {
     fn record_reducer(
         &self,
         _fingerprint: &FunctionFingerprint<'_>,
-        _energy_used: EnergyQuanta,
+        _execution_budget_used: FunctionBudget,
         _execution_duration: Duration,
     ) {
     }

crates/core/src/energy.rs

@@ -9,7 +9,7 @@ use crate::database_logger::DatabaseLogger;
 use crate::db::persistence::PersistenceProvider;
 use crate::db::relational_db::{self, spawn_view_cleanup_loop, DiskSizeFn, RelationalDB, Txdata};
 use crate::db::{self, spawn_tx_metrics_recorder};
-use crate::energy::{EnergyMonitor, EnergyQuanta, NullEnergyMonitor};
+use crate::energy::{EnergyMonitor, FunctionBudget, NullEnergyMonitor};
 use crate::host::v8::V8Runtime;
 use crate::host::ProcedureCallError;
 use crate::messages::control_db::{Database, HostType};
@@ -148,7 +148,7 @@ impl HostRuntimes {
 #[derive(Clone, Debug)]
 pub struct ReducerCallResult {
     pub outcome: ReducerOutcome,
-    pub energy_used: EnergyQuanta,
+    pub execution_budget_used: FunctionBudget,
     pub execution_duration: Duration,
 }
 

crates/core/src/host/host_controller.rs

@@ -13,7 +13,7 @@ use core::mem;
 use futures::TryFutureExt;
 use parking_lot::{Mutex, MutexGuard};
 use smallvec::SmallVec;
-use spacetimedb_client_api_messages::energy::EnergyQuanta;
+use spacetimedb_client_api_messages::energy::FunctionBudget;
 use spacetimedb_datastore::db_metrics::DB_METRICS;
 use spacetimedb_datastore::execution_context::Workload;
 use spacetimedb_datastore::locking_tx_datastore::state_view::StateView;
@@ -793,7 +793,7 @@ impl InstanceEnv {
             request_id: None,
             timer: None,
             // The procedure will pick up the tab for the energy.
-            energy_quanta_used: EnergyQuanta { quanta: 0 },
+            execution_budget_used: FunctionBudget::ZERO,
             host_execution_duration: Duration::from_millis(0),
         };
         // Commit the tx and broadcast it.

crates/core/src/host/instance_env.rs

@@ -6,7 +6,6 @@ use crate::client::messages::{OneOffQueryResponseMessage, ProcedureResultMessage
 use crate::client::{ClientActorId, ClientConnectionSender, WsVersion};
 use crate::database_logger::{DatabaseLogger, LogLevel, Record};
 use crate::db::relational_db::{RelationalDB, Tx};
-use crate::energy::EnergyQuanta;
 use crate::error::DBError;
 use crate::estimation::{check_row_limit, estimate_rows_scanned};
 use crate::hash::Hash;
@@ -212,7 +211,7 @@ pub struct ModuleEvent {
     pub function_call: ModuleFunctionCall,
     pub status: EventStatus,
     pub reducer_return_value: Option<Bytes>,
-    pub energy_quanta_used: EnergyQuanta,
+    pub execution_budget_used: FunctionBudget,
     pub host_execution_duration: Duration,
     pub request_id: Option<RequestId>,
     pub timer: Option<Instant>,
@@ -1542,7 +1541,7 @@ impl From<EventStatus> for ViewOutcome {
 pub struct ViewCallResult {
     pub outcome: ViewOutcome,
     pub tx: MutTxId,
-    pub energy_used: FunctionBudget,
+    pub execution_budget_used: FunctionBudget,
     pub total_duration: Duration,
     pub abi_duration: Duration,
 }
@@ -1551,7 +1550,7 @@ impl fmt::Debug for ViewCallResult {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("ViewCallResult")
             .field("outcome", &self.outcome)
-            .field("energy_used", &self.energy_used)
+            .field("execution_budget_used", &self.execution_budget_used)
             .field("total_duration", &self.total_duration)
             .field("abi_duration", &self.abi_duration)
             .finish()
@@ -1562,7 +1561,7 @@ impl ViewCallResult {
     pub fn default(tx: MutTxId) -> Self {
         Self {
             outcome: ViewOutcome::Success,
-            energy_used: FunctionBudget::ZERO,
+            execution_budget_used: FunctionBudget::ZERO,
             total_duration: Duration::ZERO,
             abi_duration: Duration::ZERO,
             tx,
@@ -2997,7 +2996,7 @@ impl ModuleHost {
             // Increment execution stats
             tx = result.tx;
             outcome = result.outcome;
-            energy_used += result.energy_used;
+            energy_used += result.execution_budget_used;
             total_duration += result.total_duration;
             abi_duration += result.abi_duration;
             trapped |= trap;
@@ -3011,7 +3010,7 @@ impl ModuleHost {
             ViewCallResult {
                 outcome,
                 tx,
-                energy_used,
+                execution_budget_used: energy_used,
                 total_duration,
                 abi_duration,
             },