[Darwin] Enable taking assertion on MTRDeviceController

src/darwin/Framework/CHIP/MTRDeviceController.mm

@@ -142,6 +142,10 @@ - (instancetype)initForSubclasses
         // nothing, as superclass of MTRDeviceController is NSObject
     }
     _underlyingDeviceMapLock = OS_UNFAIR_LOCK_INIT;
+
+    // Setup assertion counters
+    _keepRunningAssertionCounter = 0;
+    _shutdownPending = FALSE;
     return self;
 }
 
@@ -178,6 +182,11 @@ - (instancetype)initWithFactory:(MTRDeviceControllerFactory *)factory
         // Make sure our storage is all set up to work as early as possible,
         // before we start doing anything else with the controller.
         _uniqueIdentifier = uniqueIdentifier;
+
+        // Setup assertion counters
+        _keepRunningAssertionCounter = 0;
+        _shutdownPending = FALSE;
+
         if (storageDelegate != nil) {
             if (storageDelegateQueue == nil) {
                 MTR_LOG_ERROR("storageDelegate provided without storageDelegateQueue");
@@ -311,7 +320,53 @@ - (BOOL)isRunning
     return _cppCommissioner != nullptr;
 }
 
+- (NSUInteger)shutdownPrecheck
+{
+    __block NSUInteger assertionCount = 0;
+    dispatch_sync(_chipWorkQueue, ^{
+        self.shutdownPending = TRUE;
+        assertionCount = self.keepRunningAssertionCounter;
+    });
+    return assertionCount;
+}
+
+- (void)addRunAssertion
+{
+    dispatch_sync(_chipWorkQueue, ^{
+        ++self.keepRunningAssertionCounter;
+        MTR_LOG("%@ Adding keep running assertion, total %lu", self, (unsigned long) self.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;
+            }
+        }
+    });
+    if (shouldShutdown) {
+        MTR_LOG("%@ All assertions removed and shutdown is pending, shutting down", self);
+        [self finalShutdown];
+    }
+}
+
 - (void)shutdown
+{
+    NSUInteger assertionCount = [self shutdownPrecheck];
+    if (assertionCount != 0) {
+        MTR_LOG("%@ Pending shutdown since %lu assertions are present", self, (unsigned long) assertionCount);
+        return;
+    }
+    [self finalShutdown];
+}
+
+- (void)finalShutdown
 {
     MTR_LOG("%@ shutdown called", self);
     if (_cppCommissioner == nullptr) {
@@ -374,6 +429,7 @@ - (void)shutDownCppController
             _operationalCredentialsDelegate->SetDeviceCommissioner(nullptr);
         }
     }
+    self.shutdownPending = FALSE;
 }
 
 - (void)deinitFromFactory

src/darwin/Framework/CHIP/MTRDeviceController_Concrete.mm

@@ -326,6 +326,16 @@ - (BOOL)isRunning
 }
 
 - (void)shutdown
+{
+    NSUInteger assertionCount = [self shutdownPrecheck];
+    if (assertionCount != 0) {
+        MTR_LOG("%@ Pending shutdown since %lu assertions are present", self, (unsigned long) assertionCount);
+        return;
+    }
+    [self finalShutdown];
+}
+
+- (void)finalShutdown
 {
     MTR_LOG("%@ shutdown called", self);
     if (_cppCommissioner == nullptr) {

src/darwin/Framework/CHIP/MTRDeviceController_Internal.h

@@ -72,6 +72,10 @@ 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
@@ -289,6 +293,27 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)directlyGetSessionForNode:(chip::NodeId)nodeID completion:(MTRInternalDeviceConnectionCallback)completion;
 
+/**
+ * Takes an assertion to keep the controller running. If `-[MTRDeviceController shutdown]` is called while an assertion
+ * is held, the shutdown will be honored only after all assertions are released. Invoking this method multiple times increases
+ * the number of assertions and needs to be matched with equal amount of '-[MTRDeviceController removeRunAssertion]` to release
+ * the assertion.
+ */
+- (void)addRunAssertion;
+
+/**
+ * Removes an assertion to allow the controller to shutdown once all assertions have been released.
+ * Invoking this method once all assertions have been released in a noop.
+ */
+- (void)removeRunAssertion;
+
+/**
+ * This methods marks a request to shutdown.
+ * Returns the number of run assertions currently being held. If the value returned is not zero, it implies shutdown has to be delayed
+ * until all assertions have been removed.
+ */
+- (NSUInteger)shutdownPrecheck;
+
 @end
 
 /**