From 910c152a4e66d7ec0dfe64562e554fa73845d879 Mon Sep 17 00:00:00 2001
From: Anush Nadathur <anush@apple.com>
Date: Thu, 22 Aug 2024 06:24:17 -0700
Subject: [PATCH] [Darwin] Enable taking assertion on MTRDeviceController

- Added ability to take assertion on MTRDeviceController to keep it
  running until all assertions are removed
- A request to shutdown will take place only if there are no assertions
  are present
---
 .../Framework/CHIP/MTRDeviceController.mm     | 56 +++++++++++++++++++
 .../CHIP/MTRDeviceController_Concrete.mm      | 10 ++++
 .../CHIP/MTRDeviceController_Internal.h       | 25 +++++++++
 3 files changed, 91 insertions(+)

diff --git a/src/darwin/Framework/CHIP/MTRDeviceController.mm b/src/darwin/Framework/CHIP/MTRDeviceController.mm
index 11fd481b48b2fc..3af017da66ca8a 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceController.mm
+++ b/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, 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, 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, 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
diff --git a/src/darwin/Framework/CHIP/MTRDeviceController_Concrete.mm b/src/darwin/Framework/CHIP/MTRDeviceController_Concrete.mm
index 0ea7bf9232d766..d26d01bbc1cc0c 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceController_Concrete.mm
+++ b/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, assertionCount);
+        return;
+    }
+    [self finalShutdown];
+}
+
+- (void)finalShutdown
 {
     MTR_LOG("%@ shutdown called", self);
     if (_cppCommissioner == nullptr) {
diff --git a/src/darwin/Framework/CHIP/MTRDeviceController_Internal.h b/src/darwin/Framework/CHIP/MTRDeviceController_Internal.h
index 54d5cfd8d340fa..29a694642dad02 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceController_Internal.h
+++ b/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
 
 /**