SERVER-46548 Wait for context to update in ThreadSafetyContext unit-tests

1 files changed, 13 insertions, 3 deletions

diff --git a/src/mongo/util/thread_safety_context_test.cpp b/src/mongo/util/thread_safety_context_test.cpp
index 4e27614b514..94d3c45d5b5 100644
--- a/src/mongo/util/thread_safety_context_test.cpp
+++ b/src/mongo/util/thread_safety_context_test.cpp
@@ -29,6 +29,7 @@
 #include <vector>
 
 #include "mongo/stdx/thread.h"
+#include "mongo/unittest/barrier.h"
 #include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
 #include "mongo/util/thread_safety_context.h"
@@ -71,12 +72,21 @@ TEST_F(ThreadSafetyContextTest, CreateThreadsWithNoSafetyContext) {
 
 DEATH_TEST_F(ThreadSafetyContextTest, CreateThreadsAfterForbidingMultiThreading, "invariant") {
     ThreadSafetyContext::getThreadSafetyContext()->forbidMultiThreading();
-    auto thread = stdx::thread([] {});
-    // Must never reach here or the test fails
+    // Must terminate after starting the thread
+    auto thread = stdx::thread([] { sleepFor(Milliseconds(50)); });
+    thread.join();
 }
 
 DEATH_TEST_F(ThreadSafetyContextTest, ForbidMultiThreadingAfterCreatingThreads, "invariant") {
-    auto thread = stdx::thread([] { sleepFor(Milliseconds(50)); });
+    unittest::Barrier barrier(2);
+
+    auto thread = stdx::thread([&]() mutable {
+        barrier.countDownAndWait();
+        sleepFor(Milliseconds(50));
+    });
+
+    // Wait for the thread to start before proceeding with the test
+    barrier.countDownAndWait();
 
     ThreadSafetyContext::getThreadSafetyContext()->forbidMultiThreading();
     // Must never reach here or the test fails