From 5afe17a15d04cf16bd569ff9a9ee4699bbed81b8 Mon Sep 17 00:00:00 2001
From: Rahat Mahmood <rahat@google.com>
Date: Wed, 27 Apr 2022 14:03:59 -0700
Subject: [PATCH] cgroupfs: Fix test race with task exit.

The PID controller isn't uncharged synchronously with task exit. The
test checks the PID is released after thread join, which races with
the task being reaped.

PiperOrigin-RevId: 444960688
---
 test/syscalls/linux/cgroup.cc | 3 ++-
 test/util/cgroup_util.cc      | 8 ++++++++
 test/util/cgroup_util.h       | 5 +++++
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/test/syscalls/linux/cgroup.cc b/test/syscalls/linux/cgroup.cc
index a53c9895d..421df540e 100644
--- a/test/syscalls/linux/cgroup.cc
+++ b/test/syscalls/linux/cgroup.cc
@@ -1214,7 +1214,8 @@ TEST(PIDsCgroup, LimitEnforced) {
               IsPosixErrorOkAndHolds(baseline + 1));
 
   // Exit the first thread and try create a thread again, which should succeed.
-  t1.Join();
+  ASSERT_NO_ERRNO(child.PollControlFileForChangeAfter(
+      "pids.current", absl::Seconds(30), [&t1]() { t1.Join(); }));
   ASSERT_THAT(child.ReadIntegerControlFile("pids.current"),
               IsPosixErrorOkAndHolds(baseline));
   NoopThreads t2(1);
diff --git a/test/util/cgroup_util.cc b/test/util/cgroup_util.cc
index 4e687481f..dca02c6ce 100644
--- a/test/util/cgroup_util.cc
+++ b/test/util/cgroup_util.cc
@@ -87,6 +87,12 @@ PosixErrorOr<absl::flat_hash_set<pid_t>> Cgroup::Tasks() const {
 
 PosixError Cgroup::PollControlFileForChange(absl::string_view name,
                                             absl::Duration timeout) const {
+  return PollControlFileForChangeAfter(name, timeout, []() {});
+}
+
+PosixError Cgroup::PollControlFileForChangeAfter(
+    absl::string_view name, absl::Duration timeout,
+    std::function<void()> body) const {
   const absl::Duration poll_interval = absl::Milliseconds(10);
   const absl::Time deadline = absl::Now() + timeout;
   const std::string alias_path = absl::StrFormat("[cg#%d]/%s", id_, name);
@@ -94,6 +100,8 @@ PosixError Cgroup::PollControlFileForChange(absl::string_view name,
   ASSIGN_OR_RETURN_ERRNO(const int64_t initial_value,
                          ReadIntegerControlFile(name));
 
+  body();
+
   while (true) {
     ASSIGN_OR_RETURN_ERRNO(const int64_t current_value,
                            ReadIntegerControlFile(name));
diff --git a/test/util/cgroup_util.h b/test/util/cgroup_util.h
index f45945b3e..a13f1772d 100644
--- a/test/util/cgroup_util.h
+++ b/test/util/cgroup_util.h
@@ -78,6 +78,11 @@ class Cgroup {
   PosixError PollControlFileForChange(absl::string_view name,
                                       absl::Duration timeout) const;
 
+  // Waits for a control file's value to change after calling body.
+  PosixError PollControlFileForChangeAfter(absl::string_view name,
+                                           absl::Duration timeout,
+                                           std::function<void()> body) const;
+
   // Returns the thread ids of the leaders of thread groups managed by this
   // cgroup.
   PosixErrorOr<absl::flat_hash_set<pid_t>> Procs() const;