Clean db pool on drop only if assertions fail or it's the last iteration

crates/gpui/src/test.rs

@@ -33,6 +33,7 @@ pub fn run_test(
         Rc<platform::test::ForegroundPlatform>,
         Arc<executor::Deterministic>,
         u64,
+        bool,
     )),
 ) {
     let is_randomized = num_iterations > 1;
@@ -56,10 +57,8 @@ pub fn run_test(
             let font_cache = Arc::new(FontCache::new(font_system));
 
             loop {
-                let seed = atomic_seed.load(SeqCst);
+                let seed = atomic_seed.fetch_add(1, SeqCst);
-                if seed >= starting_seed + num_iterations {
+                let is_last_iteration = seed + 1 >= starting_seed + num_iterations;
-                    break;
-                }
 
                 if is_randomized {
                     dbg!(seed);
@@ -74,9 +73,19 @@ pub fn run_test(
                     font_cache.clone(),
                     0,
                 );
-                cx.update(|cx| test_fn(cx, foreground_platform.clone(), deterministic, seed));
+                cx.update(|cx| {
+                    test_fn(
+                        cx,
+                        foreground_platform.clone(),
+                        deterministic,
+                        seed,
+                        is_last_iteration,
+                    )
+                });
 
-                atomic_seed.fetch_add(1, SeqCst);
+                if is_last_iteration {
+                    break;
+                }
             }
         });
 

crates/gpui_macros/src/gpui_macros.rs

@@ -85,7 +85,10 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
                             ));
                         }
                         Some("StdRng") => {
-                            inner_fn_args.extend(quote!(rand::SeedableRng::seed_from_u64(seed)));
+                            inner_fn_args.extend(quote!(rand::SeedableRng::seed_from_u64(seed),));
+                        }
+                        Some("bool") => {
+                            inner_fn_args.extend(quote!(is_last_iteration,));
                         }
                         _ => {
                             return TokenStream::from(
@@ -115,7 +118,9 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
                     #num_iterations as u64,
                     #starting_seed as u64,
                     #max_retries,
-                    &mut |cx, foreground_platform, deterministic, seed| cx.foreground().run(#inner_fn_name(#inner_fn_args))
+                    &mut |cx, foreground_platform, deterministic, seed, is_last_iteration| {
+                        cx.foreground().run(#inner_fn_name(#inner_fn_args))
+                    }
                 );
             }
         }
@@ -125,9 +130,15 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
             if let FnArg::Typed(arg) = arg {
                 if let Type::Path(ty) = &*arg.ty {
                     let last_segment = ty.path.segments.last();
-                    if let Some("StdRng") = last_segment.map(|s| s.ident.to_string()).as_deref() {
+                    match last_segment.map(|s| s.ident.to_string()).as_deref() {
+                        Some("StdRng") => {
                             inner_fn_args.extend(quote!(rand::SeedableRng::seed_from_u64(seed),));
                         }
+                        Some("bool") => {
+                            inner_fn_args.extend(quote!(is_last_iteration,));
+                        }
+                        _ => {}
+                    }
                 } else {
                     inner_fn_args.extend(quote!(cx,));
                 }
@@ -147,7 +158,7 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
                     #num_iterations as u64,
                     #starting_seed as u64,
                     #max_retries,
-                    &mut |cx, _, _, seed| #inner_fn_name(#inner_fn_args)
+                    &mut |cx, _, _, seed, is_last_iteration| #inner_fn_name(#inner_fn_args)
                 );
             }
         }

crates/server/src/db.rs

@@ -537,6 +537,7 @@ pub mod tests {
         mem,
         path::Path,
         sync::atomic::{AtomicUsize, Ordering::SeqCst},
+        thread,
     };
     use util::ResultExt as _;
 
@@ -544,6 +545,7 @@ pub mod tests {
         pub db: Option<Db>,
         pub name: String,
         pub url: String,
+        clean_pool_on_drop: bool,
     }
 
     lazy_static! {
@@ -578,10 +580,15 @@ pub mod tests {
                     db: Some(db),
                     name,
                     url,
+                    clean_pool_on_drop: false,
                 }
             }
         }
 
+        pub fn set_clean_pool_on_drop(&mut self, delete_on_drop: bool) {
+            self.clean_pool_on_drop = delete_on_drop;
+        }
+
         pub fn db(&self) -> &Db {
             self.db.as_ref().unwrap()
         }
@@ -630,8 +637,11 @@ pub mod tests {
                     db: Some(db),
                     name: mem::take(&mut self.name),
                     url: mem::take(&mut self.url),
+                    clean_pool_on_drop: true,
                 });
-                if DB_COUNT.fetch_sub(1, SeqCst) == 1 {
+                if DB_COUNT.fetch_sub(1, SeqCst) == 1
+                    && (self.clean_pool_on_drop || thread::panicking())
+                {
                     block_on(async move {
                         let mut pool = DB_POOL.lock();
                         for db in pool.drain(..) {

crates/server/src/rpc.rs

@@ -1134,14 +1134,18 @@ mod tests {
     }
 
     #[gpui::test(iterations = 10)]
-    async fn test_share_project(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+    async fn test_share_project(
+        mut cx_a: TestAppContext,
+        mut cx_b: TestAppContext,
+        last_iteration: bool,
+    ) {
         let (window_b, _) = cx_b.add_window(|_| EmptyView);
         let lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
         cx_a.foreground().forbid_parking();
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -1273,13 +1277,17 @@ mod tests {
     }
 
     #[gpui::test(iterations = 10)]
-    async fn test_unshare_project(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+    async fn test_unshare_project(
+        mut cx_a: TestAppContext,
+        mut cx_b: TestAppContext,
+        last_iteration: bool,
+    ) {
         let lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
         cx_a.foreground().forbid_parking();
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -1374,13 +1382,14 @@ mod tests {
         mut cx_a: TestAppContext,
         mut cx_b: TestAppContext,
         mut cx_c: TestAppContext,
+        last_iteration: bool,
     ) {
         let lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
         cx_a.foreground().forbid_parking();
 
         // Connect to a server as 3 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
         let client_c = server.create_client(&mut cx_c, "user_c").await;
@@ -1552,13 +1561,17 @@ mod tests {
     }
 
     #[gpui::test(iterations = 10)]
-    async fn test_buffer_conflict_after_save(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+    async fn test_buffer_conflict_after_save(
+        mut cx_a: TestAppContext,
+        mut cx_b: TestAppContext,
+        last_iteration: bool,
+    ) {
         cx_a.foreground().forbid_parking();
         let lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -1640,13 +1653,17 @@ mod tests {
     }
 
     #[gpui::test(iterations = 10)]
-    async fn test_buffer_reloading(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+    async fn test_buffer_reloading(
+        mut cx_a: TestAppContext,
+        mut cx_b: TestAppContext,
+        last_iteration: bool,
+    ) {
         cx_a.foreground().forbid_parking();
         let lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -1725,13 +1742,14 @@ mod tests {
     async fn test_editing_while_guest_opens_buffer(
         mut cx_a: TestAppContext,
         mut cx_b: TestAppContext,
+        last_iteration: bool,
     ) {
         cx_a.foreground().forbid_parking();
         let lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -1807,13 +1825,14 @@ mod tests {
     async fn test_leaving_worktree_while_opening_buffer(
         mut cx_a: TestAppContext,
         mut cx_b: TestAppContext,
+        last_iteration: bool,
     ) {
         cx_a.foreground().forbid_parking();
         let lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -1882,13 +1901,17 @@ mod tests {
     }
 
     #[gpui::test(iterations = 10)]
-    async fn test_peer_disconnection(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+    async fn test_peer_disconnection(
+        mut cx_a: TestAppContext,
+        mut cx_b: TestAppContext,
+        last_iteration: bool,
+    ) {
         cx_a.foreground().forbid_parking();
         let lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -1956,6 +1979,7 @@ mod tests {
     async fn test_collaborating_with_diagnostics(
         mut cx_a: TestAppContext,
         mut cx_b: TestAppContext,
+        last_iteration: bool,
     ) {
         cx_a.foreground().forbid_parking();
         let mut lang_registry = Arc::new(LanguageRegistry::new());
@@ -1977,7 +2001,7 @@ mod tests {
             )));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -2180,6 +2204,7 @@ mod tests {
     async fn test_collaborating_with_completion(
         mut cx_a: TestAppContext,
         mut cx_b: TestAppContext,
+        last_iteration: bool,
     ) {
         cx_a.foreground().forbid_parking();
         let mut lang_registry = Arc::new(LanguageRegistry::new());
@@ -2211,7 +2236,7 @@ mod tests {
             )));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -2382,7 +2407,11 @@ mod tests {
     }
 
     #[gpui::test(iterations = 10)]
-    async fn test_formatting_buffer(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+    async fn test_formatting_buffer(
+        mut cx_a: TestAppContext,
+        mut cx_b: TestAppContext,
+        last_iteration: bool,
+    ) {
         cx_a.foreground().forbid_parking();
         let mut lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
@@ -2403,7 +2432,7 @@ mod tests {
             )));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -2484,7 +2513,11 @@ mod tests {
     }
 
     #[gpui::test(iterations = 10)]
-    async fn test_definition(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+    async fn test_definition(
+        mut cx_a: TestAppContext,
+        mut cx_b: TestAppContext,
+        last_iteration: bool,
+    ) {
         cx_a.foreground().forbid_parking();
         let mut lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
@@ -2520,7 +2553,7 @@ mod tests {
             )));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -2636,6 +2669,7 @@ mod tests {
         mut cx_a: TestAppContext,
         mut cx_b: TestAppContext,
         mut rng: StdRng,
+        last_iteration: bool,
     ) {
         cx_a.foreground().forbid_parking();
         let mut lang_registry = Arc::new(LanguageRegistry::new());
@@ -2667,7 +2701,7 @@ mod tests {
             )));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -2745,6 +2779,7 @@ mod tests {
     async fn test_collaborating_with_code_actions(
         mut cx_a: TestAppContext,
         mut cx_b: TestAppContext,
+        last_iteration: bool,
     ) {
         cx_a.foreground().forbid_parking();
         let mut lang_registry = Arc::new(LanguageRegistry::new());
@@ -2774,7 +2809,7 @@ mod tests {
             )));
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -2983,11 +3018,15 @@ mod tests {
     }
 
     #[gpui::test(iterations = 10)]
-    async fn test_basic_chat(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+    async fn test_basic_chat(
+        mut cx_a: TestAppContext,
+        mut cx_b: TestAppContext,
+        last_iteration: bool,
+    ) {
         cx_a.foreground().forbid_parking();
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
 
@@ -3123,10 +3162,10 @@ mod tests {
     }
 
     #[gpui::test(iterations = 10)]
-    async fn test_chat_message_validation(mut cx_a: TestAppContext) {
+    async fn test_chat_message_validation(mut cx_a: TestAppContext, last_iteration: bool) {
         cx_a.foreground().forbid_parking();
 
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
 
         let db = &server.app_state.db;
@@ -3183,11 +3222,15 @@ mod tests {
     }
 
     #[gpui::test(iterations = 10)]
-    async fn test_chat_reconnection(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+    async fn test_chat_reconnection(
+        mut cx_a: TestAppContext,
+        mut cx_b: TestAppContext,
+        last_iteration: bool,
+    ) {
         cx_a.foreground().forbid_parking();
 
         // Connect to a server as 2 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
         let mut status_b = client_b.status();
@@ -3399,13 +3442,14 @@ mod tests {
         mut cx_a: TestAppContext,
         mut cx_b: TestAppContext,
         mut cx_c: TestAppContext,
+        last_iteration: bool,
     ) {
         cx_a.foreground().forbid_parking();
         let lang_registry = Arc::new(LanguageRegistry::new());
         let fs = Arc::new(FakeFs::new(cx_a.background()));
 
         // Connect to a server as 3 clients.
-        let mut server = TestServer::start(cx_a.foreground()).await;
+        let mut server = TestServer::start(cx_a.foreground(), last_iteration).await;
         let client_a = server.create_client(&mut cx_a, "user_a").await;
         let client_b = server.create_client(&mut cx_b, "user_b").await;
         let client_c = server.create_client(&mut cx_c, "user_c").await;
@@ -3536,8 +3580,8 @@ mod tests {
         }
     }
 
-    #[gpui::test(iterations = 10)]
+    #[gpui::test(iterations = 100)]
-    async fn test_random_collaboration(cx: TestAppContext, rng: StdRng) {
+    async fn test_random_collaboration(cx: TestAppContext, rng: StdRng, last_iteration: bool) {
         cx.foreground().forbid_parking();
         let max_peers = env::var("MAX_PEERS")
             .map(|i| i.parse().expect("invalid `MAX_PEERS` variable"))
@@ -3558,7 +3602,7 @@ mod tests {
         .await;
 
         let operations = Rc::new(Cell::new(0));
-        let mut server = TestServer::start(cx.foreground()).await;
+        let mut server = TestServer::start(cx.foreground(), last_iteration).await;
         let mut clients = Vec::new();
 
         let mut next_entity_id = 100000;
@@ -3729,8 +3773,9 @@ mod tests {
     }
 
     impl TestServer {
-        async fn start(foreground: Rc<executor::Foreground>) -> Self {
+        async fn start(foreground: Rc<executor::Foreground>, clean_db_pool_on_drop: bool) -> Self {
-            let test_db = TestDb::new();
+            let mut test_db = TestDb::new();
+            test_db.set_clean_pool_on_drop(clean_db_pool_on_drop);
             let app_state = Self::build_app_state(&test_db).await;
             let peer = Peer::new();
             let notifications = mpsc::channel(128);