From e1adbf0217df3c4bda09ee3c2466265406d75539 Mon Sep 17 00:00:00 2001
From: Christian Kamm <mail@ckamm.de>
Date: Sat, 27 Aug 2022 17:42:21 +0200
Subject: [PATCH] Tests: Fix concurrent log capture

---
 programs/mango-v4/tests/program_test/mod.rs   | 19 ++++++++------
 .../mango-v4/tests/program_test/solana.rs     | 26 +++++++++++++++----
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/programs/mango-v4/tests/program_test/mod.rs b/programs/mango-v4/tests/program_test/mod.rs
index 50f3a2cf1..2ca8772b3 100644
--- a/programs/mango-v4/tests/program_test/mod.rs
+++ b/programs/mango-v4/tests/program_test/mod.rs
@@ -52,7 +52,7 @@ impl AddPacked for ProgramTest {
 
 struct LoggerWrapper {
     inner: env_logger::Logger,
-    program_log: Arc<RwLock<Vec<String>>>,
+    capture: Arc<RwLock<Vec<String>>>,
 }
 
 impl Log for LoggerWrapper {
@@ -67,9 +67,9 @@ impl Log for LoggerWrapper {
         {
             let msg = record.args().to_string();
             if let Some(data) = msg.strip_prefix("Program log: ") {
-                self.program_log.write().unwrap().push(data.into());
+                self.capture.write().unwrap().push(data.into());
             } else if let Some(data) = msg.strip_prefix("Program data: ") {
-                self.program_log.write().unwrap().push(data.into());
+                self.capture.write().unwrap().push(data.into());
             }
         }
         self.inner.log(record);
@@ -87,12 +87,13 @@ pub struct MarginTradeCookie {
 
 pub struct TestContextBuilder {
     test: ProgramTest,
-    program_log_capture: Arc<RwLock<Vec<String>>>,
+    logger_capture: Arc<RwLock<Vec<String>>>,
     mint0: Pubkey,
 }
 
 lazy_static::lazy_static! {
-    static ref PROGRAM_LOG_CAPTURE: Arc<RwLock<Vec<String>>> = Arc::new(RwLock::new(vec![]));
+    static ref LOGGER_CAPTURE: Arc<RwLock<Vec<String>>> = Arc::new(RwLock::new(vec![]));
+    static ref LOGGER_LOCK: Arc<RwLock<()>> = Arc::new(RwLock::new(()));
 }
 
 impl TestContextBuilder {
@@ -108,7 +109,7 @@ impl TestContextBuilder {
                 .build();
         let _ = log::set_boxed_logger(Box::new(LoggerWrapper {
             inner: env_logger,
-            program_log: PROGRAM_LOG_CAPTURE.clone(),
+            capture: LOGGER_CAPTURE.clone(),
         }));
 
         let mut test = ProgramTest::new("mango_v4", mango_v4::id(), processor!(mango_v4::entry));
@@ -118,7 +119,7 @@ impl TestContextBuilder {
 
         Self {
             test,
-            program_log_capture: PROGRAM_LOG_CAPTURE.clone(),
+            logger_capture: LOGGER_CAPTURE.clone(),
             mint0: Pubkey::new_unique(),
         }
     }
@@ -282,7 +283,9 @@ impl TestContextBuilder {
         let solana = Arc::new(SolanaCookie {
             context: RefCell::new(context),
             rent,
-            program_log: self.program_log_capture.clone(),
+            logger_capture: self.logger_capture.clone(),
+            logger_lock: LOGGER_LOCK.clone(),
+            last_transaction_log: RefCell::new(vec![]),
         });
 
         solana
diff --git a/programs/mango-v4/tests/program_test/solana.rs b/programs/mango-v4/tests/program_test/solana.rs
index 170cf102c..ac0195bca 100644
--- a/programs/mango-v4/tests/program_test/solana.rs
+++ b/programs/mango-v4/tests/program_test/solana.rs
@@ -17,7 +17,9 @@ use spl_token::*;
 pub struct SolanaCookie {
     pub context: RefCell<ProgramTestContext>,
     pub rent: Rent,
-    pub program_log: Arc<RwLock<Vec<String>>>,
+    pub logger_capture: Arc<RwLock<Vec<String>>>,
+    pub logger_lock: Arc<RwLock<()>>,
+    pub last_transaction_log: RefCell<Vec<String>>,
 }
 
 impl SolanaCookie {
@@ -27,7 +29,15 @@ impl SolanaCookie {
         instructions: &[Instruction],
         signers: Option<&[&Keypair]>,
     ) -> Result<(), BanksClientError> {
-        self.program_log.write().unwrap().clear();
+        // The locking in this function is convoluted:
+        // We capture the program log output by overriding the global logger and capturing
+        // messages there. This logger is potentially shared among multiple tests that run
+        // concurrently.
+        // To allow each independent SolanaCookie to capture only the logs from the transaction
+        // passed to process_transaction, wo globally hold the "program_log_lock" for the
+        // duration that the tx needs to process. So only a single one can run at a time.
+        let tx_log_lock = Arc::new(self.logger_lock.write().unwrap());
+        self.logger_capture.write().unwrap().clear();
 
         let mut context = self.context.borrow_mut();
 
@@ -45,13 +55,19 @@ impl SolanaCookie {
 
         transaction.sign(&all_signers, context.last_blockhash);
 
-        context
+        let result = context
             .banks_client
             .process_transaction_with_commitment(
                 transaction,
                 solana_sdk::commitment_config::CommitmentLevel::Processed,
             )
-            .await
+            .await;
+
+        *self.last_transaction_log.borrow_mut() = self.logger_capture.read().unwrap().clone();
+
+        drop(tx_log_lock);
+
+        result
     }
 
     pub async fn get_clock(&self) -> solana_program::clock::Clock {
@@ -221,7 +237,7 @@ impl SolanaCookie {
 
     #[allow(dead_code)]
     pub fn program_log(&self) -> Vec<String> {
-        self.program_log.read().unwrap().clone()
+        self.last_transaction_log.borrow().clone()
     }
 
     pub fn program_log_events<T: anchor_lang::Event + anchor_lang::AnchorDeserialize>(