From e2ed35bbc0a11ceb3dd00680a17f5026379f7d75 Mon Sep 17 00:00:00 2001
From: Anush Nadathur <anush@apple.com>
Date: Thu, 22 Aug 2024 18:06:10 -0700
Subject: [PATCH] Simplified to use lock for assertion tracking

---
 .../Framework/CHIP/MTRDeviceController.mm     | 77 +++++++++----------
 .../CHIP/MTRDeviceController_Internal.h       |  4 -
 2 files changed, 38 insertions(+), 43 deletions(-)

diff --git a/src/darwin/Framework/CHIP/MTRDeviceController.mm b/src/darwin/Framework/CHIP/MTRDeviceController.mm
index b7c2e1395dd6ac..d87a6c3c210326 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceController.mm
+++ b/src/darwin/Framework/CHIP/MTRDeviceController.mm
@@ -136,6 +136,11 @@ @implementation MTRDeviceController {
     NSNumber * _fabricID;
     NSNumber * _nodeID;
     NSData * _rootPublicKey;
+
+    // Counters to track assertion status and access controlled by the _assertionLock
+    NSUInteger _keepRunningAssertionCounter;
+    BOOL _shutdownPending;
+    os_unfair_lock _assertionLock;
 }
 
 - (os_unfair_lock_t)deviceMapLock
@@ -150,9 +155,10 @@ - (instancetype)initForSubclasses
     }
     _underlyingDeviceMapLock = OS_UNFAIR_LOCK_INIT;
 
-    // Setup assertion counters
+    // Setup assertion variables
     _keepRunningAssertionCounter = 0;
-    _shutdownPending = FALSE;
+    _shutdownPending = NO;
+    _assertionLock = OS_UNFAIR_LOCK_INIT;
     return self;
 }
 
@@ -190,9 +196,10 @@ - (instancetype)initWithFactory:(MTRDeviceControllerFactory *)factory
         // before we start doing anything else with the controller.
         _uniqueIdentifier = uniqueIdentifier;
 
-        // Setup assertion counters
+        // Setup assertion variables
         _keepRunningAssertionCounter = 0;
-        _shutdownPending = FALSE;
+        _shutdownPending = NO;
+        _assertionLock = OS_UNFAIR_LOCK_INIT;
 
         if (storageDelegate != nil) {
             if (storageDelegateQueue == nil) {
@@ -330,16 +337,6 @@ - (BOOL)isRunning
     return _cppCommissioner != nullptr;
 }
 
-- (NSUInteger)shutdownPrecheck
-{
-    __block NSUInteger assertionCount = 0;
-    dispatch_sync(_chipWorkQueue, ^{
-        self.shutdownPending = TRUE;
-        assertionCount = self.keepRunningAssertionCounter;
-    });
-    return assertionCount;
-}
-
 - (BOOL)matchesPendingShutdownWithParams:(MTRDeviceControllerParameters *)parameters
 {
     if (!parameters.operationalCertificate || !parameters.rootCertificate) {
@@ -351,8 +348,8 @@ - (BOOL)matchesPendingShutdownWithParams:(MTRDeviceControllerParameters *)parame
 
     __block BOOL matches = FALSE;
     dispatch_sync(_chipWorkQueue, ^{
-        matches =  self.keepRunningAssertionCounter > 0 &&
-                   self.shutdownPending &&
+        matches =  _keepRunningAssertionCounter > 0 &&
+                   _shutdownPending &&
                    MTREqualObjects(nodeID, self->_nodeID) &&
                    MTREqualObjects(fabricID, self->_fabricID) &&
                    MTREqualObjects(publicKey, self->_rootPublicKey);
@@ -362,42 +359,42 @@ - (BOOL)matchesPendingShutdownWithParams:(MTRDeviceControllerParameters *)parame
 
 - (void)addRunAssertion
 {
-    dispatch_sync(_chipWorkQueue, ^{
-        ++self.keepRunningAssertionCounter;
-        MTR_LOG("%@ Adding keep running assertion, total %lu", self, (unsigned long) self.keepRunningAssertionCounter);
-    });
+    std::lock_guard lock(_assertionLock);
+
+    // Only take an assertion if running
+    if ([self isRunning]) {
+        ++_keepRunningAssertionCounter;
+        MTR_LOG("%@ Adding keep running assertion, total %lu", self, (unsigned long) _keepRunningAssertionCounter);
+    }
 }
 
 - (void)removeRunAssertion;
 {
-    __block BOOL shouldShutdown = FALSE;
-    dispatch_sync(_chipWorkQueue, ^{
-        if (self.keepRunningAssertionCounter > 0) {
-            --self.keepRunningAssertionCounter;
-            MTR_LOG("%@ Removing keep running assertion, total %lu", self, (unsigned long) self.keepRunningAssertionCounter);
-            if (self.keepRunningAssertionCounter == 0 && self.shutdownPending) {
-                shouldShutdown = TRUE;
-            }
+    std::lock_guard lock(_assertionLock);
+
+    if (_keepRunningAssertionCounter > 0) {
+        --_keepRunningAssertionCounter;
+        MTR_LOG("%@ Removing keep running assertion, total %lu", self, (unsigned long) _keepRunningAssertionCounter);
+
+        if ([self isRunning] && _keepRunningAssertionCounter == 0 && _shutdownPending) {
+            MTR_LOG("%@ All assertions removed and shutdown is pending, shutting down", self);
+            [self finalShutdown];
         }
-    });
-    if (shouldShutdown) {
-        MTR_LOG("%@ All assertions removed and shutdown is pending, shutting down", self);
-        [self finalShutdown];
     }
 }
 
 - (void)clearPendingShutdown
 {
-    dispatch_sync(_chipWorkQueue, ^{
-        self.shutdownPending = FALSE;
-    });
+    std::lock_guard lock(_assertionLock);
+    _shutdownPending = NO;
 }
 
 - (void)shutdown
 {
-    NSUInteger assertionCount = [self shutdownPrecheck];
-    if (assertionCount != 0) {
-        MTR_LOG("%@ Pending shutdown since %lu assertions are present", self, (unsigned long) assertionCount);
+    std::lock_guard lock(_assertionLock);
+
+    if (_keepRunningAssertionCounter > 0) {
+        MTR_LOG("%@ Pending shutdown since %lu assertions are present", self, (unsigned long) _keepRunningAssertionCounter);
         return;
     }
     [self finalShutdown];
@@ -405,6 +402,8 @@ - (void)shutdown
 
 - (void)finalShutdown
 {
+    os_unfair_lock_assert_owner(&_assertionLock);
+
     MTR_LOG("%@ shutdown called", self);
     if (_cppCommissioner == nullptr) {
         // Already shut down.
@@ -469,7 +468,7 @@ - (void)shutDownCppController
             _operationalCredentialsDelegate->SetDeviceCommissioner(nullptr);
         }
     }
-    self.shutdownPending = FALSE;
+    _shutdownPending = NO;
 }
 
 - (void)deinitFromFactory
diff --git a/src/darwin/Framework/CHIP/MTRDeviceController_Internal.h b/src/darwin/Framework/CHIP/MTRDeviceController_Internal.h
index 5c4caa49cdd00f..37392c02bb5eda 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceController_Internal.h
+++ b/src/darwin/Framework/CHIP/MTRDeviceController_Internal.h
@@ -73,10 +73,6 @@ NS_ASSUME_NONNULL_BEGIN
 // (moved here so subclasses can initialize differently)
 @property (readwrite, retain) dispatch_queue_t chipWorkQueue;
 
-// Counters to track assertion status
-@property (nonatomic, readwrite) NSUInteger keepRunningAssertionCounter;
-@property (nonatomic, readwrite) BOOL shutdownPending;
-
 - (instancetype)initForSubclasses;
 
 #pragma mark - MTRDeviceControllerFactory methods