PidWaiter: Add a thread-safe Notify

PiperOrigin-RevId: 721291791
Change-Id: I115155e01c6f916f6881108a5cae9053887fdf79

Author: happyCoder92

sandboxed_api/sandbox2/util/BUILD

@@ -159,6 +159,9 @@ cc_library(
     copts = sapi_platform_copts(),
     deps = [
         ":deadline_manager",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/cleanup",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
     ],
 )
@@ -171,5 +174,6 @@ cc_test(
         ":pid_waiter",
         "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest_main",
+        "@com_google_sandboxed_api//sandboxed_api/util:thread",
     ],
 )

sandboxed_api/sandbox2/util/CMakeLists.txt

@@ -68,7 +68,10 @@ add_library(sandbox2_util_pid_waiter ${SAPI_LIB_TYPE}
 )
 add_library(sandbox2::pid_waiter ALIAS sandbox2_util_pid_waiter)
 target_link_libraries(sandbox2_util_pid_waiter
-  PUBLIC absl::time
+  PUBLIC absl::core_headers
+         absl::cleanup
+         absl::synchronization
+         absl::time
          sandbox2::deadline_manager
   PRIVATE sapi::base
 )
@@ -166,6 +169,7 @@ if(BUILD_TESTING AND SAPI_BUILD_TESTING)
     absl::time
     sandbox2::pid_waiter
     sapi::test_main
+    sapi::thread
   )
 
   # sandboxed_api/sandbox2/util:deadline_manager_test

sandboxed_api/sandbox2/util/pid_waiter.cc

@@ -19,6 +19,11 @@
 #include <cerrno>
 #include <memory>
 
+#include "absl/cleanup/cleanup.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/time/time.h"
+#include "sandboxed_api/sandbox2/util/deadline_manager.h"
+
 namespace sandbox2 {
 
 namespace {
@@ -75,6 +80,10 @@ bool PidWaiter::CheckStatus(pid_t pid, bool blocking) {
 void PidWaiter::RefillStatuses() {
   constexpr int kMaxIterations = 1000;
   constexpr int kPriorityCheckPeriod = 100;
+  absl::Cleanup notify = [this] {
+    absl::MutexLock lock(&notify_mutex_);
+    notified_ = false;
+  };
   if (!statuses_.empty()) {
     return;
   }
@@ -88,8 +97,22 @@ void PidWaiter::RefillStatuses() {
       break;
     }
   }
-  if (statuses_.empty() && deadline_registration_.has_value()) {
-    deadline_registration_->ExecuteBlockingSyscall(
+  if (statuses_.empty()) {
+    DeadlineRegistration* deadline_registration = nullptr;
+    {
+      absl::MutexLock lock(&notify_mutex_);
+      if (deadline_ == absl::InfinitePast() || notified_) {
+        return;
+      }
+      if (!deadline_registration_.has_value()) {
+        deadline_registration_.emplace(DeadlineManager::instance());
+      }
+      deadline_registration_->SetDeadline(deadline_);
+      // DeadlineRegistration is only constructed once, so it's safe to use it
+      // outside of the lock.
+      deadline_registration = &*deadline_registration_;
+    }
+    deadline_registration->ExecuteBlockingSyscall(
         [&] { CheckStatus(-1, /*blocking=*/true); });
   }
 }

sandboxed_api/sandbox2/util/pid_waiter.h

@@ -22,6 +22,8 @@
 #include <optional>
 #include <utility>
 
+#include "absl/base/thread_annotations.h"
+#include "absl/synchronization/mutex.h"
 #include "absl/time/time.h"
 #include "sandboxed_api/sandbox2/util/deadline_manager.h"
 
@@ -58,21 +60,34 @@ class PidWaiter {
 
   // Sets the deadline for the next Wait() call.
   void SetDeadline(absl::Time deadline) {
-    if (!deadline_registration_.has_value()) {
-      deadline_registration_.emplace(DeadlineManager::instance());
+    absl::MutexLock lock(&notify_mutex_);
+    deadline_ = deadline;
+  }
+
+  // Breaks out of the current Wait operation, if there is one running
+  // concurrently. Otherwise, makes the next Wait non-blocking.
+  // Can be called concurrently with Wait and SetDeadline.
+  void Notify() {
+    absl::MutexLock lock(&notify_mutex_);
+    if (deadline_registration_.has_value()) {
+      deadline_registration_->SetDeadline(absl::InfinitePast());
     }
-    deadline_registration_->SetDeadline(deadline);
+    notified_ = true;
   }
 
  private:
   bool CheckStatus(pid_t pid, bool blocking = false);
   void RefillStatuses();
 
   pid_t priority_pid_;
-  std::optional<DeadlineRegistration> deadline_registration_;
   std::deque<std::pair<pid_t, int>> statuses_;
   std::unique_ptr<WaitPidInterface> wait_pid_iface_;
   int last_errno_ = 0;
+  absl::Mutex notify_mutex_;
+  absl::Time deadline_ ABSL_GUARDED_BY(notify_mutex_) = absl::InfinitePast();
+  std::optional<DeadlineRegistration> deadline_registration_
+      ABSL_GUARDED_BY(notify_mutex_);
+  bool notified_ ABSL_GUARDED_BY(notify_mutex_) = false;
 };
 
 }  // namespace sandbox2

sandboxed_api/sandbox2/util/pid_waiter_test.cc

@@ -23,6 +23,7 @@
 #include "gtest/gtest.h"
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
+#include "sandboxed_api/util/thread.h"
 
 namespace sandbox2 {
 namespace {
@@ -153,5 +154,72 @@ TEST(PidWaiterTest, DeadlineRespected) {
   EXPECT_EQ(errno, EINTR);
 }
 
+TEST(PidWaiterTest, NotifyConcurrent) {
+  auto mock_wait_pid = std::make_unique<MockWaitPid>();
+  EXPECT_CALL(*mock_wait_pid,
+              WaitPid(_, _, __WNOTHREAD | __WALL | WUNTRACED | WNOHANG))
+      .WillRepeatedly(Return(0));
+  EXPECT_CALL(*mock_wait_pid, WaitPid(_, _, __WNOTHREAD | __WALL | WUNTRACED))
+      .WillRepeatedly([](int pid, int* status, int flags) {
+        struct timespec ts = absl::ToTimespec(absl::Seconds(2));
+        if (nanosleep(&ts, nullptr) == -1) {
+          return -1;
+        }
+        *status = kFirstPid;
+        return kFirstPid;
+      });
+  PidWaiter waiter(kPrioPid, std::move(mock_wait_pid));
+  waiter.SetDeadline(absl::Now() + absl::Seconds(1));
+  sapi::Thread notify_thread([&] {
+    absl::SleepFor(absl::Milliseconds(100));
+    waiter.Notify();
+  });
+  int status;
+  absl::Time start = absl::Now();
+  EXPECT_EQ(waiter.Wait(&status), -1);
+  EXPECT_LT(absl::Now() - start, absl::Milliseconds(500));
+  EXPECT_EQ(errno, EINTR);
+}
+
+TEST(PidWaiterTest, NotifyNext) {
+  auto mock_wait_pid = std::make_unique<MockWaitPid>();
+  EXPECT_CALL(*mock_wait_pid,
+              WaitPid(_, _, __WNOTHREAD | __WALL | WUNTRACED | WNOHANG))
+      .WillRepeatedly(Return(0));
+  EXPECT_CALL(*mock_wait_pid, WaitPid(_, _, __WNOTHREAD | __WALL | WUNTRACED))
+      .Times(0);
+  PidWaiter waiter(kPrioPid, std::move(mock_wait_pid));
+  waiter.SetDeadline(absl::Now() + absl::Seconds(1));
+  waiter.Notify();
+  int status;
+  absl::Time start = absl::Now();
+  EXPECT_EQ(waiter.Wait(&status), 0);
+  EXPECT_LT(absl::Now() - start, absl::Milliseconds(500));
+}
+
+TEST(PidWaiterTest, DeadlineUnchangedAfterNotify) {
+  auto mock_wait_pid = std::make_unique<MockWaitPid>();
+  EXPECT_CALL(*mock_wait_pid,
+              WaitPid(_, _, __WNOTHREAD | __WALL | WUNTRACED | WNOHANG))
+      .WillRepeatedly(Return(0));
+  EXPECT_CALL(*mock_wait_pid, WaitPid(_, _, __WNOTHREAD | __WALL | WUNTRACED))
+      .WillRepeatedly([](int pid, int* status, int flags) {
+        struct timespec ts = absl::ToTimespec(absl::Milliseconds(500));
+        if (nanosleep(&ts, nullptr) == -1) {
+          return -1;
+        }
+        *status = kFirstPid;
+        return kFirstPid;
+      });
+  PidWaiter waiter(kPrioPid, std::move(mock_wait_pid));
+  waiter.SetDeadline(absl::Now() + absl::Milliseconds(900));
+  waiter.Notify();
+  int status;
+  EXPECT_EQ(waiter.Wait(&status), 0);
+  absl::SleepFor(absl::Milliseconds(500));
+  EXPECT_EQ(waiter.Wait(&status), -1);
+  EXPECT_EQ(errno, EINTR);
+}
+
 }  // namespace
 }  // namespace sandbox2