From a3851689458a11d456722a61ea93c5080c0598d6 Mon Sep 17 00:00:00 2001
From: Jeff Tung <100387939+jtung-apple@users.noreply.github.com>
Date: Tue, 9 Jul 2024 14:13:06 -0700
Subject: [PATCH 1/2] [Darwin] MTRDeviceController should always examine the
 MTRDevice before using subscription/CASE pool

---
 src/darwin/Framework/CHIP/MTRDevice.mm         |  3 +++
 .../Framework/CHIP/MTRDeviceController.mm      | 18 ++++++++++++------
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/darwin/Framework/CHIP/MTRDevice.mm b/src/darwin/Framework/CHIP/MTRDevice.mm
index 30d96d3c2e9f43..2491855c5775e5 100644
--- a/src/darwin/Framework/CHIP/MTRDevice.mm
+++ b/src/darwin/Framework/CHIP/MTRDevice.mm
@@ -539,6 +539,9 @@ - (instancetype)initWithNodeID:(NSNumber *)nodeID controller:(MTRDeviceControlle
 - (void)dealloc
 {
     [[NSNotificationCenter defaultCenter] removeObserver:_systemTimeChangeObserverToken];
+
+    // TODO: retain cycle and clean up https://github.com/project-chip/connectedhomeip/issues/34267
+    MTR_LOG("%@ dealloc called", self);
 }
 
 - (NSString *)description
diff --git a/src/darwin/Framework/CHIP/MTRDeviceController.mm b/src/darwin/Framework/CHIP/MTRDeviceController.mm
index cea281551384d0..f9cf8c8a4a080f 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceController.mm
+++ b/src/darwin/Framework/CHIP/MTRDeviceController.mm
@@ -605,11 +605,20 @@ - (BOOL)startup:(MTRDeviceControllerStartupParamsInternal *)startupParams
             MTR_LOG("Loaded attribute values for %lu nodes from storage for controller uuid %@", static_cast<unsigned long>(clusterDataByNode.count), self->_uniqueIdentifier);
 
             std::lock_guard lock(self->_deviceMapLock);
+            NSMutableArray * deviceList = [NSMutableArray array];
             for (NSNumber * nodeID in clusterDataByNode) {
                 NSDictionary * clusterData = clusterDataByNode[nodeID];
                 MTRDevice * device = [self _setupDeviceForNodeID:nodeID prefetchedClusterData:clusterData];
                 MTR_LOG("Loaded %lu cluster data from storage for %@", static_cast<unsigned long>(clusterData.count), device);
+
+                [deviceList addObject:device];
             }
+
+#define kSecondsToWaitBeforeAPIClientRetainsMTRDevice 60
+            // Keep the devices retained for a while, in case API client doesn't immedieately retain them
+            dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t) (kSecondsToWaitBeforeAPIClientRetainsMTRDevice * NSEC_PER_SEC)), dispatch_get_main_queue(), ^{
+                MTR_LOG("MTRDeviceController: un-retain devices loaded at startup %lu", static_cast<unsigned long>(deviceList.count));
+            });
         }];
     }
 
@@ -1257,15 +1266,12 @@ - (BOOL)checkIsRunning:(NSError * __autoreleasing *)error
 
 - (void)getSessionForNode:(chip::NodeId)nodeID completion:(MTRInternalDeviceConnectionCallback)completion
 {
-    // First check if MTRDevice exists from having loaded from storage, or created by a client.
-    // Do not use deviceForNodeID here, because we don't want to create the device if it does not already exist.
-    os_unfair_lock_lock(&_deviceMapLock);
-    MTRDevice * device = [_nodeIDToDeviceMap objectForKey:@(nodeID)];
-    os_unfair_lock_unlock(&_deviceMapLock);
+    // Get the corresponding MTRDevice object to determine if the case/subscription pool is to be used
+    MTRDevice * device = [self deviceForNodeID:@(nodeID)];
 
     // In the case that this device is known to use thread, queue this with subscription attempts as well, to
     // help with throttling Thread traffic.
-    if (device && [device deviceUsesThread]) {
+    if ([device deviceUsesThread]) {
         MTRAsyncWorkItem * workItem = [[MTRAsyncWorkItem alloc] initWithQueue:dispatch_get_global_queue(QOS_CLASS_DEFAULT, 0)];
         [workItem setReadyHandler:^(id _Nonnull context, NSInteger retryCount, MTRAsyncWorkCompletionBlock _Nonnull workItemCompletion) {
             MTRInternalDeviceConnectionCallback completionWrapper = ^(chip::Messaging::ExchangeManager * _Nullable exchangeManager,

From 0685340cf1ba5bc2548e4d26818f97f77b217b77 Mon Sep 17 00:00:00 2001
From: Jeff Tung <100387939+jtung-apple@users.noreply.github.com>
Date: Tue, 9 Jul 2024 15:58:44 -0700
Subject: [PATCH 2/2] Update src/darwin/Framework/CHIP/MTRDevice.mm

Co-authored-by: Justin Wood <woody@apple.com>
---
 src/darwin/Framework/CHIP/MTRDevice.mm | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/darwin/Framework/CHIP/MTRDevice.mm b/src/darwin/Framework/CHIP/MTRDevice.mm
index 2491855c5775e5..394364e588eac1 100644
--- a/src/darwin/Framework/CHIP/MTRDevice.mm
+++ b/src/darwin/Framework/CHIP/MTRDevice.mm
@@ -541,7 +541,7 @@ - (void)dealloc
     [[NSNotificationCenter defaultCenter] removeObserver:_systemTimeChangeObserverToken];
 
     // TODO: retain cycle and clean up https://github.com/project-chip/connectedhomeip/issues/34267
-    MTR_LOG("%@ dealloc called", self);
+    MTR_LOG("MTRDevice dealloc: %p", self);
 }
 
 - (NSString *)description