1 files changed, 9 insertions, 6 deletions
diff --git a/tests/leak_test.cpp b/tests/leak_test.cpp
index 600520959..3cc1a0a91 100644
--- a/tests/leak_test.cpp
+++ b/tests/leak_test.cpp
@@ -17,6 +17,7 @@
 #include <err.h>
 #include <inttypes.h>
 #include <pthread.h>
+#include <sched.h>
 #include <stdio.h>
 #include <string.h>
 #include <sys/mman.h>
@@ -30,26 +31,28 @@
 #include <vector>
 
 #include <android-base/macros.h>
+#include <android-base/threads.h>
 
 #include "utils.h"
 
 using namespace std::chrono_literals;
 
-static void WaitUntilAllExited(pid_t* pids, size_t pid_count) {
+static void WaitUntilAllThreadsExited(pid_t* tids, size_t tid_count) {
   // Wait until all children have exited.
   bool alive = true;
   while (alive) {
     alive = false;
-    for (size_t i = 0; i < pid_count; ++i) {
-      if (pids[i] != 0) {
-        if (kill(pids[i], 0) == 0) {
+    for (size_t i = 0; i < tid_count; ++i) {
+      if (tids[i] != 0) {
+        if (tgkill(getpid(), tids[i], 0) == 0) {
           alive = true;
         } else {
           EXPECT_EQ(errno, ESRCH);
-          pids[i] = 0;  // Skip in next loop.
+          tids[i] = 0;  // Skip in next loop.
         }
       }
     }
+    sched_yield();
   }
 }
 
@@ -155,7 +158,7 @@ TEST(pthread_leak, detach) {
     pthread_barrier_wait(&barrier);
     ASSERT_EQ(pthread_barrier_destroy(&barrier), 0);
 
-    WaitUntilAllExited(tids, arraysize(tids));
+    WaitUntilAllThreadsExited(tids, threads_count);
 
     // A native bridge implementation might need a warm up pass to reach a steady state.
     // http://b/37920774.