2nd attempt at trace configuration

Author: buffalojoec

harness/src/lib.rs

@@ -506,8 +506,11 @@ pub struct Mollusk {
     #[cfg(feature = "invocation-inspect-callback")]
     pub invocation_inspect_callback: Box<dyn InvocationInspectCallback>,
 
-    #[cfg(feature = "register-tracing")]
-    pub enable_register_tracing: bool,
+    /// Dictates whether or not register tracing was enabled.
+    /// Provided as input to the invocation inspect callback for potential
+    /// register trace consumption.
+    #[cfg(feature = "invocation-inspect-callback")]
+    enable_register_tracing: bool,
 
     /// This field stores the slot only to be able to convert to and from FD
     /// fixtures and a Mollusk instance, since FD fixtures have a
@@ -528,7 +531,7 @@ pub trait InvocationInspectCallback {
         invoke_context: &InvokeContext,
     );
 
-    fn after_invocation(&self, invoke_context: &InvokeContext);
+    fn after_invocation(&self, invoke_context: &InvokeContext, register_tracing_enabled: bool);
 }
 
 #[cfg(feature = "invocation-inspect-callback")]
@@ -539,71 +542,12 @@ impl InvocationInspectCallback for EmptyInvocationInspectCallback {
     fn before_invocation(&self, _: &Pubkey, _: &[u8], _: &[InstructionAccount], _: &InvokeContext) {
     }
 
-    fn after_invocation(&self, _: &InvokeContext) {}
+    fn after_invocation(&self, _: &InvokeContext, _register_tracing_enabled: bool) {}
 }
 
 impl Default for Mollusk {
     fn default() -> Self {
-        #[rustfmt::skip]
-        solana_logger::setup_with_default(
-            "solana_rbpf::vm=debug,\
-             solana_runtime::message_processor=debug,\
-             solana_runtime::system_instruction_processor=trace",
-        );
-        let compute_budget = ComputeBudget::new_with_defaults(true, true);
-
-        #[cfg(feature = "fuzz")]
-        let feature_set = {
-            // Omit "test features" (they have the same u64 ID).
-            let mut fs = FeatureSet::all_enabled();
-            fs.active_mut()
-                .remove(&agave_feature_set::disable_sbpf_v0_execution::id());
-            fs.active_mut()
-                .remove(&agave_feature_set::reenable_sbpf_v0_execution::id());
-            fs
-        };
-        #[cfg(not(feature = "fuzz"))]
-        let feature_set = FeatureSet::all_enabled();
-
-        let program_cache = ProgramCache::new(
-            &feature_set,
-            &compute_budget,
-            #[cfg(feature = "register-tracing")]
-            {
-                true
-            },
-            #[cfg(not(feature = "register-tracing"))]
-            {
-                false
-            },
-        );
-
-        Self {
-            config: Config::default(),
-            compute_budget,
-            epoch_stake: EpochStake::default(),
-            feature_set,
-            logger: None,
-            program_cache,
-            sysvars: Sysvars::default(),
-
-            // Register tracing feature requires `invocation-inspect-callback`.
-            // Use tracing callback when both are active.
-            #[cfg(all(feature = "invocation-inspect-callback", feature = "register-tracing"))]
-            invocation_inspect_callback: Box::new(DefaultRegisterTracingCallback::default()),
-            // Use empty callback when only `invocation-inspect-callback` is active.
-            #[cfg(all(
-                feature = "invocation-inspect-callback",
-                not(feature = "register-tracing"),
-            ))]
-            invocation_inspect_callback: Box::new(EmptyInvocationInspectCallback {}),
-
-            #[cfg(feature = "register-tracing")]
-            enable_register_tracing: true,
-
-            #[cfg(feature = "fuzz-fd")]
-            slot: 0,
-        }
+        Self::new_inner(/* enable_register_tracing */ false)
     }
 }
 
@@ -669,6 +613,61 @@ impl InvokeContextCallback for MolluskInvokeContextCallback<'_> {
 }
 
 impl Mollusk {
+    fn new_inner(#[allow(unused)] enable_register_tracing: bool) -> Self {
+        #[rustfmt::skip]
+        solana_logger::setup_with_default(
+            "solana_rbpf::vm=debug,\
+             solana_runtime::message_processor=debug,\
+             solana_runtime::system_instruction_processor=trace",
+        );
+        let compute_budget = ComputeBudget::new_with_defaults(true, true);
+
+        #[cfg(feature = "fuzz")]
+        let feature_set = {
+            // Omit "test features" (they have the same u64 ID).
+            let mut fs = FeatureSet::all_enabled();
+            fs.active_mut()
+                .remove(&agave_feature_set::disable_sbpf_v0_execution::id());
+            fs.active_mut()
+                .remove(&agave_feature_set::reenable_sbpf_v0_execution::id());
+            fs
+        };
+        #[cfg(not(feature = "fuzz"))]
+        let feature_set = FeatureSet::all_enabled();
+
+        let program_cache =
+            ProgramCache::new(&feature_set, &compute_budget, enable_register_tracing);
+
+        #[allow(unused_mut)]
+        let mut me = Self {
+            config: Config::default(),
+            compute_budget,
+            epoch_stake: EpochStake::default(),
+            feature_set,
+            logger: None,
+            program_cache,
+            sysvars: Sysvars::default(),
+
+            #[cfg(feature = "invocation-inspect-callback")]
+            invocation_inspect_callback: Box::new(EmptyInvocationInspectCallback {}),
+
+            #[cfg(feature = "invocation-inspect-callback")]
+            enable_register_tracing,
+
+            #[cfg(feature = "fuzz-fd")]
+            slot: 0,
+        };
+
+        #[cfg(feature = "register-tracing")]
+        if enable_register_tracing {
+            // Have a default register tracing callback if register tracing is
+            // enabled.
+            me.invocation_inspect_callback = Box::new(DefaultRegisterTracingCallback::default());
+        }
+
+        me
+    }
+
     /// Create a new Mollusk instance containing the provided program.
     ///
     /// Attempts to load the program's ELF file from the default search paths.
@@ -689,6 +688,28 @@ impl Mollusk {
         mollusk
     }
 
+    /// Create a new Mollusk instance with configurable debugging features.
+    ///
+    /// This constructor allows enabling low-level VM debugging capabilities,
+    /// such as register tracing, which are baked into program executables at
+    /// load time and cannot be changed afterwards.
+    ///
+    /// When `enable_register_tracing` is `true`:
+    /// - Programs are compiled with register tracing support
+    /// - A default [`DefaultRegisterTracingCallback`] is installed
+    /// - Trace data is written to `SBF_TRACE_DIR` (or `target/sbf/trace` by
+    ///   default)
+    #[cfg(feature = "register-tracing")]
+    pub fn new_debuggable(
+        program_id: &Pubkey,
+        program_name: &str,
+        enable_register_tracing: bool,
+    ) -> Self {
+        let mut mollusk = Self::new_inner(enable_register_tracing);
+        mollusk.add_program(program_id, program_name, &DEFAULT_LOADER_KEY);
+        mollusk
+    }
+
     /// Add a program to the test environment.
     ///
     /// If you intend to CPI to a program, this is likely what you want to use.
@@ -822,7 +843,7 @@ impl Mollusk {
 
             #[cfg(feature = "invocation-inspect-callback")]
             self.invocation_inspect_callback
-                .after_invocation(&invoke_context);
+                .after_invocation(&invoke_context, self.enable_register_tracing);
 
             result
         };

harness/src/register_tracing.rs

@@ -82,16 +82,19 @@ impl InvocationInspectCallback for DefaultRegisterTracingCallback {
     fn before_invocation(&self, _: &Pubkey, _: &[u8], _: &[InstructionAccount], _: &InvokeContext) {
     }
 
-    fn after_invocation(&self, invoke_context: &InvokeContext) {
-        invoke_context.iterate_vm_traces(
-            &|instruction_context: InstructionContext,
-              executable: &Executable,
-              register_trace: RegisterTrace| {
-                if let Err(e) = self.handler(instruction_context, executable, register_trace) {
-                    eprintln!("Error collecting the register tracing: {}", e);
-                }
-            },
-        );
+    fn after_invocation(&self, invoke_context: &InvokeContext, register_tracing_enabled: bool) {
+        if register_tracing_enabled {
+            // Only read the register traces if they were actually enabled.
+            invoke_context.iterate_vm_traces(
+                &|instruction_context: InstructionContext,
+                  executable: &Executable,
+                  register_trace: RegisterTrace| {
+                    if let Err(e) = self.handler(instruction_context, executable, register_trace) {
+                        eprintln!("Error collecting the register tracing: {}", e);
+                    }
+                },
+            );
+        }
     }
 }
 

harness/tests/fd_test_vectors.rs

@@ -62,12 +62,10 @@ fn test_load_firedancer_fixtures() {
                         },
                         result,
                     ) = load_firedancer_fixture(&loaded_fixture);
-                    let mollusk = Mollusk {
-                        compute_budget,
-                        feature_set,
-                        slot,
-                        ..Default::default()
-                    };
+                    let mut mollusk = Mollusk::default();
+                    mollusk.compute_budget = compute_budget;
+                    mollusk.feature_set = feature_set;
+                    mollusk.slot = slot;
                     let generated_fixture =
                         build_fixture_from_mollusk_test(&mollusk, &instruction, &accounts, &result);
 

harness/tests/register_tracing.rs

@@ -69,24 +69,34 @@ fn test_custom_register_tracing_callback() {
         ) {
         }
 
-        fn after_invocation(&self, invoke_context: &InvokeContext) {
-            invoke_context.iterate_vm_traces(
-                &|instruction_context: InstructionContext,
-                  executable: &Executable,
-                  register_trace: RegisterTrace| {
-                    if let Err(e) = self.handler(instruction_context, executable, register_trace) {
-                        eprintln!("Error collecting the register tracing: {}", e);
-                    }
-                },
-            );
+        fn after_invocation(&self, invoke_context: &InvokeContext, register_tracing_enabled: bool) {
+            // Only process traces if register tracing was enabled.
+            if register_tracing_enabled {
+                invoke_context.iterate_vm_traces(
+                    &|instruction_context: InstructionContext,
+                      executable: &Executable,
+                      register_trace: RegisterTrace| {
+                        if let Err(e) =
+                            self.handler(instruction_context, executable, register_trace)
+                        {
+                            eprintln!("Error collecting the register tracing: {}", e);
+                        }
+                    },
+                );
+            }
         }
     }
 
     std::env::set_var("SBF_OUT_DIR", "../target/deploy");
 
     let program_id = Pubkey::new_unique();
     let payer_pk = Pubkey::new_unique();
-    let mut mollusk = Mollusk::new(&program_id, "test_program_primary");
+    // Use new_debuggable with register tracing enabled.
+    let mut mollusk = Mollusk::new_debuggable(
+        &program_id,
+        "test_program_primary",
+        /* enable_register_tracing */ true,
+    );
 
     // Phase 1 - basic register tracing test.
 
@@ -136,31 +146,49 @@ fn test_custom_register_tracing_callback() {
             collected_data.executed_jump_instructions_count;
     }
 
-    // Phase 2 - check we can stop register tracing for this instance of Mollusk.
+    // Phase 2 - check that register tracing is disabled when constructing
+    // Mollusk with enable_register_tracing=false.
     {
         // Clear the tracing data collected so far.
         {
             let mut td = tracing_data.borrow_mut();
             td.clear();
         }
 
-        mollusk.enable_register_tracing = false;
+        // Create a new Mollusk instance with register tracing disabled.
+        let mut mollusk_no_tracing = Mollusk::new_debuggable(
+            &program_id,
+            "test_program_primary",
+            /* enable_register_tracing */ false,
+        );
+        mollusk_no_tracing.invocation_inspect_callback = Box::new(CustomRegisterTracingCallback {
+            tracing_data: Rc::clone(&tracing_data),
+        });
 
         // Execute the same instruction again.
-        let _ = mollusk.process_instruction(&instruction, &accounts);
+        let _ = mollusk_no_tracing.process_instruction(&instruction, &accounts);
 
         let td = tracing_data.borrow();
-        // We expect it to be empty!
+        // We expect it to be empty since tracing was disabled!
         assert!(td.is_empty());
     }
 
-    // Phase 3 - check we can have register tracing back for this instance of
+    // Phase 3 - check we can have register tracing enabled for a new instance of
     // Mollusk.
     {
-        mollusk.enable_register_tracing = true;
+        // Create a new Mollusk instance with register tracing enabled.
+        let mut mollusk_with_tracing = Mollusk::new_debuggable(
+            &program_id,
+            "test_program_primary",
+            /* enable_register_tracing */ true,
+        );
+        mollusk_with_tracing.invocation_inspect_callback =
+            Box::new(CustomRegisterTracingCallback {
+                tracing_data: Rc::clone(&tracing_data),
+            });
 
         // Execute the same instruction again.
-        let _ = mollusk.process_instruction(&instruction, &accounts);
+        let _ = mollusk_with_tracing.process_instruction(&instruction, &accounts);
 
         let td = tracing_data.borrow();
         let collected_data = td.get(&program_id).unwrap();