wasm: refactor energy; v8: energy, timings, and timeout

Author: Centril

crates/core/src/host/v8/mod.rs

@@ -10,20 +10,23 @@ use crate::host::wasm_common::module_host_actor::{
 use crate::host::ArgsTuple;
 use crate::{host::Scheduler, module_host_context::ModuleCreationContext, replica_context::ReplicaContext};
 use anyhow::anyhow;
-use std::time::Instant;
+use core::sync::atomic::{AtomicBool, Ordering};
+use core::time::Duration;
 use de::deserialize_js;
 use error::{catch_exception, exception_already_thrown, log_traceback, ExcResult, Throwable};
 use from_value::cast;
 use key_cache::get_or_create_key_cache;
 use ser::serialize_to_js;
-use spacetimedb_client_api_messages::energy::{EnergyQuanta, ReducerBudget};
+use spacetimedb_client_api_messages::energy::ReducerBudget;
 use spacetimedb_datastore::locking_tx_datastore::MutTxId;
 use spacetimedb_datastore::traits::Program;
 use spacetimedb_lib::RawModuleDef;
 use spacetimedb_lib::{ConnectionId, Identity};
 use spacetimedb_schema::auto_migrate::MigrationPolicy;
 use std::sync::{Arc, LazyLock};
-use v8::{Context, ContextOptions, ContextScope, Function, HandleScope, Isolate, Local, Value};
+use std::thread;
+use std::time::Instant;
+use v8::{Context, ContextOptions, ContextScope, Function, HandleScope, Isolate, IsolateHandle, Local, Value};
 
 mod de;
 mod error;
@@ -152,25 +155,28 @@ impl ModuleInstance for JsInstance {
             tx,
             params,
             log_traceback,
-            |tx, op, _budget| {
+            |tx, op, budget| {
                 // TODO(centril): snapshots, module->host calls
                 // Setup V8 scope.
                 let mut isolate: v8::OwnedIsolate = Isolate::new(<_>::default());
+                let isolate_handle = isolate.thread_safe_handle();
                 let mut scope_1 = HandleScope::new(&mut isolate);
                 let context = Context::new(&mut scope_1, ContextOptions::default());
                 let mut scope_2 = ContextScope::new(&mut scope_1, context);
 
+                let timeout_thread_cancel_flag = run_reducer_timeout(isolate_handle, budget);
+
                 // Call the reducer.
                 let start = Instant::now();
                 let call_result = call_call_reducer_from_op(&mut scope_2, op);
                 let total_duration = start.elapsed();
+                // Cancel the execution timeout in `run_reducer_timeout`.
+                timeout_thread_cancel_flag.store(true, Ordering::Relaxed);
 
-                // TODO(centril): energy metrering.
-                let energy = EnergyStats {
-                    used: EnergyQuanta::ZERO,
-                    wasmtime_fuel_used: 0,
-                    remaining: ReducerBudget::ZERO,
-                };
+                // Handle energy and timings.
+                let used = duration_to_budget(total_duration);
+                let remaining = budget - used;
+                let energy = EnergyStats { budget, remaining };
                 let timings = ExecutionTimings {
                     total_duration,
                     // TODO(centril): call times.
@@ -195,6 +201,41 @@ impl ModuleInstance for JsInstance {
     }
 }
 
+fn run_reducer_timeout(isolate_handle: IsolateHandle, budget: ReducerBudget) -> Arc<AtomicBool> {
+    let execution_done_flag = Arc::new(AtomicBool::new(false));
+    let execution_done_flag2 = execution_done_flag.clone();
+    let timeout = budget_to_duration(budget);
+
+    // TODO(centril): Using an OS thread is a bit heavy handed...?
+    thread::spawn(move || {
+        // Sleep until the timeout.
+        thread::sleep(timeout);
+
+        if execution_done_flag2.load(Ordering::Relaxed) {
+            // The reducer completed successfully.
+            return;
+        }
+
+        // Reducer is still running.
+        // Terminate V8 execution.
+        isolate_handle.terminate_execution();
+    });
+
+    execution_done_flag
+}
+
+fn budget_to_duration(_budget: ReducerBudget) -> Duration {
+    // TODO(centril): This is fake logic that allows a maximum timeout.
+    // Replace with sensible math.
+    Duration::MAX
+}
+
+fn duration_to_budget(_duration: Duration) -> ReducerBudget {
+    // TODO(centril): This is fake logic that allows minimum energy usage.
+    // Replace with sensible math.
+    ReducerBudget::ZERO
+}
+
 /// Returns the global property `key`.
 fn get_global_property<'scope>(
     scope: &mut HandleScope<'scope>,

crates/core/src/host/wasm_common/module_host_actor.rs

@@ -8,7 +8,7 @@ use tracing::span::EnteredSpan;
 use super::instrumentation::CallTimes;
 use crate::client::ClientConnectionSender;
 use crate::database_logger;
-use crate::energy::{EnergyMonitor, EnergyQuanta, ReducerBudget, ReducerFingerprint};
+use crate::energy::{EnergyMonitor, ReducerBudget, ReducerFingerprint};
 use crate::host::instance_env::InstanceEnv;
 use crate::host::module_common::{build_common_module_from_raw, ModuleCommon};
 use crate::host::module_host::{
@@ -60,11 +60,17 @@ pub trait WasmInstance: Send + Sync + 'static {
 }
 
 pub struct EnergyStats {
-    pub used: EnergyQuanta,
-    pub wasmtime_fuel_used: u64,
+    pub budget: ReducerBudget,
     pub remaining: ReducerBudget,
 }
 
+impl EnergyStats {
+    /// Returns the used energy amount.
+    fn used(&self) -> ReducerBudget {
+        (self.budget.get() - self.remaining.get()).into()
+    }
+}
+
 pub struct ExecutionTimings {
     pub total_duration: Duration,
     pub wasm_instance_env_call_times: CallTimes,
@@ -412,22 +418,24 @@ impl InstanceCommon {
             call_result,
         } = result;
 
+        let energy_used = energy.used();
+        let energy_quanta_used = energy_used.into();
         vm_metrics.report(
-            energy.wasmtime_fuel_used,
+            energy_used.get(),
             timings.total_duration,
             &timings.wasm_instance_env_call_times,
         );
 
         self.energy_monitor
-            .record_reducer(&energy_fingerprint, energy.used, timings.total_duration);
+            .record_reducer(&energy_fingerprint, energy_quanta_used, timings.total_duration);
         if self.allocated_memory != memory_allocation {
             self.metric_wasm_memory_bytes.set(memory_allocation as i64);
             self.allocated_memory = memory_allocation;
         }
 
         reducer_span
             .record("timings.total_duration", tracing::field::debug(timings.total_duration))
-            .record("energy.used", tracing::field::debug(energy.used));
+            .record("energy.used", tracing::field::debug(energy_used));
 
         maybe_log_long_running_reducer(reducer_name, timings.total_duration);
         reducer_span.exit();
@@ -486,7 +494,7 @@ impl InstanceCommon {
                 args,
             },
             status,
-            energy_quanta_used: energy.used,
+            energy_quanta_used,
             host_execution_duration: timings.total_duration,
             request_id,
             timer,
@@ -495,7 +503,7 @@ impl InstanceCommon {
 
         ReducerCallResult {
             outcome: ReducerOutcome::from(&event.status),
-            energy_used: energy.used,
+            energy_used: energy_quanta_used,
             execution_duration: timings.total_duration,
         }
     }

crates/core/src/host/wasmtime/wasmtime_module.rs

@@ -192,11 +192,8 @@ impl module_host_actor::WasmInstance for WasmtimeInstance {
     #[tracing::instrument(level = "trace", skip_all)]
     fn call_reducer(&mut self, op: ReducerOp<'_>, budget: ReducerBudget) -> module_host_actor::ExecuteResult {
         let store = &mut self.store;
-        // note that ReducerBudget being a u64 is load-bearing here - although we convert budget right back into
-        // EnergyQuanta at the end of this function, from_energy_quanta clamps it to a u64 range.
-        // otherwise, we'd return something like `used: i128::MAX - u64::MAX`, which is inaccurate.
+        // Set the fuel budget in WASM.
         set_store_fuel(store, budget.into());
-        let original_fuel = get_store_fuel(store);
         store.set_epoch_deadline(EPOCH_TICKS_PER_SECOND);
 
         // Prepare sender identity and connection ID, as LITTLE-ENDIAN byte arrays.
@@ -231,14 +228,10 @@ impl module_host_actor::WasmInstance for WasmtimeInstance {
 
         let call_result = call_result.map(|code| handle_error_sink_code(code, error));
 
+        // Compute fuel and heap usage.
         let remaining_fuel = get_store_fuel(store);
-
         let remaining: ReducerBudget = remaining_fuel.into();
-        let energy = module_host_actor::EnergyStats {
-            used: (budget - remaining).into(),
-            wasmtime_fuel_used: original_fuel.0 - remaining_fuel.0,
-            remaining,
-        };
+        let energy = module_host_actor::EnergyStats { budget, remaining };
         let memory_allocation = store.data().get_mem().memory.data_size(&store);
 
         module_host_actor::ExecuteResult {