Add internal interface to determine if a SpinLock is cooperative.

Some use cases of SpinLock need to verify that it is configured for
non-cooperative scheduling.

PiperOrigin-RevId: 575278400
Change-Id: Ic49f57a989a3f7f67e346a1ec545f4cd5b94f849

2 files changed, 26 insertions, 0 deletions

diff --git a/absl/base/internal/spinlock.h b/absl/base/internal/spinlock.h
index 09ba5824..6818a86a 100644
--- a/absl/base/internal/spinlock.h
+++ b/absl/base/internal/spinlock.h
@@ -41,6 +41,14 @@
 #include "absl/base/internal/tsan_mutex_interface.h"
 #include "absl/base/thread_annotations.h"
 
+namespace tcmalloc {
+namespace tcmalloc_internal {
+
+class AllocationGuardSpinLockHolder;
+
+}  // namespace tcmalloc_internal
+}  // namespace tcmalloc
+
 namespace absl {
 ABSL_NAMESPACE_BEGIN
 namespace base_internal {
@@ -137,6 +145,7 @@ class ABSL_LOCKABLE SpinLock {
 
   // Provide access to protected method above.  Use for testing only.
   friend struct SpinLockTest;
+  friend class tcmalloc::tcmalloc_internal::AllocationGuardSpinLockHolder;
 
  private:
   // lockword_ is used to store the following:
@@ -171,6 +180,10 @@ class ABSL_LOCKABLE SpinLock {
     return scheduling_mode == base_internal::SCHEDULE_COOPERATIVE_AND_KERNEL;
   }
 
+  bool IsCooperative() const {
+    return lockword_.load(std::memory_order_relaxed) & kSpinLockCooperative;
+  }
+
   uint32_t TryLockInternal(uint32_t lock_value, uint32_t wait_cycles);
   void SlowLock() ABSL_ATTRIBUTE_COLD;
   void SlowUnlock(uint32_t lock_value) ABSL_ATTRIBUTE_COLD;
diff --git a/absl/base/spinlock_test_common.cc b/absl/base/spinlock_test_common.cc
index 52ecf580..e9047158 100644
--- a/absl/base/spinlock_test_common.cc
+++ b/absl/base/spinlock_test_common.cc
@@ -51,6 +51,8 @@ struct SpinLockTest {
   static int64_t DecodeWaitCycles(uint32_t lock_value) {
     return SpinLock::DecodeWaitCycles(lock_value);
   }
+
+  static bool IsCooperative(const SpinLock& l) { return l.IsCooperative(); }
 };
 
 namespace {
@@ -266,6 +268,17 @@ TEST(SpinLockWithThreads, DoesNotDeadlock) {
                        base_internal::NumCPUs() * 2);
 }
 
+TEST(SpinLockTest, IsCooperative) {
+  SpinLock default_constructor;
+  EXPECT_TRUE(SpinLockTest::IsCooperative(default_constructor));
+
+  SpinLock cooperative(base_internal::SCHEDULE_COOPERATIVE_AND_KERNEL);
+  EXPECT_TRUE(SpinLockTest::IsCooperative(cooperative));
+
+  SpinLock kernel_only(base_internal::SCHEDULE_KERNEL_ONLY);
+  EXPECT_FALSE(SpinLockTest::IsCooperative(kernel_only));
+}
+
 }  // namespace
 }  // namespace base_internal
 ABSL_NAMESPACE_END