From 9a469fe30519a052abaacf167a3d52df79ba67f5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= <AlexanderMeissner@gmx.net>
Date: Fri, 10 Feb 2023 11:30:00 +0100
Subject: [PATCH] Makes undeployment visible immediately.

---
 program-runtime/src/executor_cache.rs |  6 +++++-
 programs/bpf_loader/src/lib.rs        |  9 +++++++++
 programs/sbf/tests/programs.rs        | 10 +++-------
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/program-runtime/src/executor_cache.rs b/program-runtime/src/executor_cache.rs
index 9a3aacd92f7580..22b68cb90a829c 100644
--- a/program-runtime/src/executor_cache.rs
+++ b/program-runtime/src/executor_cache.rs
@@ -43,6 +43,10 @@ impl TransactionExecutorCache {
         self.visible.get(key).cloned()
     }
 
+    pub fn set_tombstone(&mut self, key: Pubkey) {
+        self.visible.insert(key, None);
+    }
+
     pub fn set(
         &mut self,
         key: Pubkey,
@@ -54,7 +58,7 @@ impl TransactionExecutorCache {
             if delay_visibility_of_program_deployment {
                 // Place a tombstone in the cache so that
                 // we don't load the new version from the database as it should remain invisible
-                self.visible.insert(key, None);
+                self.set_tombstone(key);
             } else {
                 self.visible.insert(key, Some(executor.clone()));
             }
diff --git a/programs/bpf_loader/src/lib.rs b/programs/bpf_loader/src/lib.rs
index 459e2dfc315f6a..462961c536f19c 100644
--- a/programs/bpf_loader/src/lib.rs
+++ b/programs/bpf_loader/src/lib.rs
@@ -1166,6 +1166,15 @@ fn process_loader_upgradeable_instruction(
                                 instruction_context,
                                 &log_collector,
                             )?;
+                            if invoke_context
+                                .feature_set
+                                .is_active(&delay_visibility_of_program_deployment::id())
+                            {
+                                invoke_context
+                                    .tx_executor_cache
+                                    .borrow_mut()
+                                    .set_tombstone(program_key);
+                            }
                         }
                         _ => {
                             ic_logger_msg!(log_collector, "Invalid Program account");
diff --git a/programs/sbf/tests/programs.rs b/programs/sbf/tests/programs.rs
index 6a60d04bef6606..4961c351ff60c6 100644
--- a/programs/sbf/tests/programs.rs
+++ b/programs/sbf/tests/programs.rs
@@ -2031,7 +2031,7 @@ fn test_program_sbf_invoke_in_same_tx_as_undeployment() {
         Some(&program_id),
     );
 
-    // Undeployment is invisible to both top-level-instructions and CPI instructions
+    // Undeployment is visible to both top-level-instructions and CPI instructions
     for invoke_instruction in [invoke_instruction, indirect_invoke_instruction] {
         // Call upgradeable program
         let result =
@@ -2040,11 +2040,7 @@ fn test_program_sbf_invoke_in_same_tx_as_undeployment() {
 
         // Upgrade the program and invoke in same tx
         let message = Message::new(
-            &[
-                undeployment_instruction.clone(),
-                invoke_instruction,
-                panic_instruction.clone(), // Make sure the TX fails, so we don't have to deploy another program
-            ],
+            &[undeployment_instruction.clone(), invoke_instruction],
             Some(&mint_keypair.pubkey()),
         );
         let tx = Transaction::new(
@@ -2055,7 +2051,7 @@ fn test_program_sbf_invoke_in_same_tx_as_undeployment() {
         let (result, _, _) = process_transaction_and_record_inner(&bank, tx);
         assert_eq!(
             result.unwrap_err(),
-            TransactionError::InstructionError(2, InstructionError::ProgramFailedToComplete),
+            TransactionError::InstructionError(1, InstructionError::InvalidAccountData),
         );
     }
 }