Add MTRDeviceController API to get a list of node IDs with data stored for them. (#37344)

* Add MTRDeviceController API to get a list of node IDs with data stored for them.

To avoid races between modifications to _nodesWithResumptionInfo and our reading
it, added a lock to protect _nodesWithResumptionInfo.  Also added a cached
_nodesWithAttributeInfo (protected by the same lock) so that we don't have to do
a sync dispatch to our storage queue to return a value from this API.

* Fix unit test that was relying on the storage being the source of truth for the node index.

Author: bzbarsky-apple

src/darwin/Framework/CHIP/MTRDeviceController.h

@@ -309,6 +309,13 @@ MTR_AVAILABLE(ios(16.1), macos(13.0), watchos(9.1), tvos(16.1))
  */
 - (void)resume MTR_AVAILABLE(ios(18.2), macos(15.2), watchos(11.2), tvos(18.2));
 
+/**
+ * Returns the list of node IDs for which this controller has stored
+ * information.  Returns empty list if the controller does not have any
+ * information stored.
+ */
+- (NSArray<NSNumber *> *)nodesWithStoredData MTR_AVAILABLE(ios(18.4), macos(15.4), watchos(11.4), tvos(18.4));
+
 /**
  * Shut down the controller. Calls to shutdown after the first one are NO-OPs.
  * This must be called, either directly or via shutting down the

src/darwin/Framework/CHIP/MTRDeviceController.mm

@@ -508,6 +508,12 @@ + (void)forceLocalhostAdvertisingOnly
 }
 #endif // DEBUG
 
+- (NSArray<NSNumber *> *)nodesWithStoredData
+{
+    MTR_ABSTRACT_METHOD();
+    return @[];
+}
+
 #pragma mark - MTRDeviceControllerDelegate management
 
 // Note these are implemented in the base class so that XPC subclass can use it as well

src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.h

@@ -105,6 +105,12 @@ typedef void (^MTRDeviceControllerDataStoreClusterDataHandler)(NSDictionary<NSNu
  */
 - (void)synchronouslyPerformBlock:(void (^_Nullable)(void))block;
 
+/**
+ * Returns the list of node IDs for which this data store has stored data of
+ * some sort.
+ */
+- (NSArray<NSNumber *> *)nodesWithStoredData;
+
 @end
 
 NS_ASSUME_NONNULL_END

src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm

@@ -19,6 +19,7 @@
 // Importing MTRBaseDevice.h for the MTRAttributePath class. Needs to change when https://github.com/project-chip/connectedhomeip/issues/31247 is fixed.
 #import "MTRBaseDevice.h"
 #import "MTRLogging_Internal.h"
+#import "MTRUnfairLock.h"
 
 #include <lib/core/CASEAuthTag.h>
 #include <lib/core/NodeId.h>
@@ -111,6 +112,11 @@ @implementation MTRDeviceControllerDataStore {
     __weak MTRDeviceController * _controller;
     // Array of nodes with resumption info, oldest-stored first.
     NSMutableArray<NSNumber *> * _nodesWithResumptionInfo;
+    // Array of nodes with attribute info.
+    NSMutableArray<NSNumber *> * _nodesWithAttributeInfo;
+    // Lock protecting access to the _nodesWithAttributeInfo and
+    // _nodesWithResumptionInfo arrays.
+    os_unfair_lock _nodeArrayLock;
 }
 
 - (nullable instancetype)initWithController:(MTRDeviceController *)controller
@@ -124,8 +130,10 @@ - (nullable instancetype)initWithController:(MTRDeviceController *)controller
     _controller = controller;
     _storageDelegate = storageDelegate;
     _storageDelegateQueue = storageDelegateQueue;
+    _nodeArrayLock = OS_UNFAIR_LOCK_INIT;
 
     __block id resumptionNodeList;
+    __block NSArray<NSNumber *> * nodesWithAttributeInfo;
     dispatch_sync(_storageDelegateQueue, ^{
         @autoreleasepool {
             // NOTE: controller, not our weak ref, since we know it's still
@@ -134,6 +142,8 @@ - (nullable instancetype)initWithController:(MTRDeviceController *)controller
                                                   valueForKey:sResumptionNodeListKey
                                                 securityLevel:MTRStorageSecurityLevelSecure
                                                   sharingType:MTRStorageSharingTypeNotShared];
+
+            nodesWithAttributeInfo = [self _fetchNodeIndex];
         }
     });
     if (resumptionNodeList != nil) {
@@ -152,6 +162,12 @@ - (nullable instancetype)initWithController:(MTRDeviceController *)controller
         _nodesWithResumptionInfo = [[NSMutableArray alloc] init];
     }
 
+    if (nodesWithAttributeInfo != nil) {
+        _nodesWithAttributeInfo = [nodesWithAttributeInfo mutableCopy];
+    } else {
+        _nodesWithAttributeInfo = [[NSMutableArray alloc] init];
+    }
+
     return self;
 }
 
@@ -196,6 +212,8 @@ - (void)storeResumptionInfo:(MTRCASESessionResumptionInfo *)resumptionInfo
                        removeValueForKey:ResumptionByResumptionIDKey(oldInfo.resumptionID)
                            securityLevel:MTRStorageSecurityLevelSecure
                              sharingType:MTRStorageSharingTypeNotShared];
+
+            std::lock_guard lock(self->_nodeArrayLock);
             [_nodesWithResumptionInfo removeObject:resumptionInfo.nodeID];
         }
 
@@ -211,9 +229,14 @@ - (void)storeResumptionInfo:(MTRCASESessionResumptionInfo *)resumptionInfo
                          sharingType:MTRStorageSharingTypeNotShared];
 
         // Update our resumption info node list.
-        [_nodesWithResumptionInfo addObject:resumptionInfo.nodeID];
+        NSArray<NSNumber *> * valueToStore;
+        {
+            std::lock_guard lock(self->_nodeArrayLock);
+            [_nodesWithResumptionInfo addObject:resumptionInfo.nodeID];
+            valueToStore = [_nodesWithResumptionInfo copy];
+        }
         [_storageDelegate controller:controller
-                          storeValue:[_nodesWithResumptionInfo copy]
+                          storeValue:valueToStore
                               forKey:sResumptionNodeListKey
                        securityLevel:MTRStorageSecurityLevelSecure
                          sharingType:MTRStorageSharingTypeNotShared];
@@ -226,7 +249,9 @@ - (void)clearAllResumptionInfo
     VerifyOrReturn(controller != nil); // No way to call delegate without controller.
 
     // Can we do less dispatch?  We would need to have a version of
-    // _findResumptionInfoWithKey that assumes we are already on the right queue.
+    // _findResumptionInfoWithKey that assumes we are already on the right
+    // queue.
+    std::lock_guard lock(_nodeArrayLock);
     for (NSNumber * nodeID in _nodesWithResumptionInfo) {
         [self _clearResumptionInfoForNodeID:nodeID controller:controller];
     }
@@ -240,6 +265,8 @@ - (void)clearResumptionInfoForNodeID:(NSNumber *)nodeID
     VerifyOrReturn(controller != nil); // No way to call delegate without controller.
 
     [self _clearResumptionInfoForNodeID:nodeID controller:controller];
+
+    std::lock_guard lock(_nodeArrayLock);
     [_nodesWithResumptionInfo removeObject:nodeID];
 }
 
@@ -588,6 +615,20 @@ - (void)unitTestPruneEmptyStoredClusterDataBranches
         [self _pruneEmptyStoredClusterDataBranches];
     });
 }
+
+- (void)unitTestRereadNodeIndex
+{
+    dispatch_sync(_storageDelegateQueue, ^{
+        auto * newIndex = [self _fetchNodeIndex];
+
+        std::lock_guard lock(self->_nodeArrayLock);
+        if (newIndex != nil) {
+            self->_nodesWithAttributeInfo = [newIndex mutableCopy];
+        } else {
+            self->_nodesWithAttributeInfo = [[NSMutableArray alloc] init];
+        }
+    });
+}
 #endif
 
 - (void)_pruneEmptyStoredClusterDataBranches
@@ -596,9 +637,8 @@ - (void)_pruneEmptyStoredClusterDataBranches
 
     NSUInteger storeFailures = 0;
 
-    // Fetch node index
-    NSArray<NSNumber *> * nodeIndex = [self _fetchNodeIndex];
-    NSMutableArray<NSNumber *> * nodeIndexCopy = [nodeIndex mutableCopy];
+    std::lock_guard lock(self->_nodeArrayLock);
+    NSArray<NSNumber *> * nodeIndex = [_nodesWithAttributeInfo copy];
 
     for (NSNumber * nodeID in nodeIndex) {
         // Fetch endpoint index
@@ -638,7 +678,7 @@ - (void)_pruneEmptyStoredClusterDataBranches
             if (endpointIndexCopy.count) {
                 success = [self _storeEndpointIndex:endpointIndexCopy forNodeID:nodeID];
             } else {
-                [nodeIndexCopy removeObject:nodeID];
+                [_nodesWithAttributeInfo removeObject:nodeID];
                 success = [self _deleteEndpointIndexForNodeID:nodeID];
             }
             if (!success) {
@@ -648,16 +688,16 @@ - (void)_pruneEmptyStoredClusterDataBranches
         }
     }
 
-    if (nodeIndex.count != nodeIndexCopy.count) {
+    if (nodeIndex.count != _nodesWithAttributeInfo.count) {
         BOOL success;
-        if (nodeIndexCopy.count) {
-            success = [self _storeNodeIndex:nodeIndexCopy];
+        if (_nodesWithAttributeInfo.count) {
+            success = [self _storeNodeIndex:_nodesWithAttributeInfo];
         } else {
             success = [self _deleteNodeIndex];
         }
         if (!success) {
             storeFailures++;
-            MTR_LOG_ERROR("Store failed in _pruneEmptyStoredClusterDataBranches for nodeIndex (%lu)", static_cast<unsigned long>(nodeIndexCopy.count));
+            MTR_LOG_ERROR("Store failed in _pruneEmptyStoredClusterDataBranches for nodeIndex (%lu)", static_cast<unsigned long>(_nodesWithAttributeInfo.count));
         }
     }
 
@@ -713,18 +753,19 @@ - (void)clearStoredClusterDataForNodeID:(NSNumber *)nodeID
 {
     dispatch_async(_storageDelegateQueue, ^{
         [self _clearStoredClusterDataForNodeID:nodeID];
-        NSArray<NSNumber *> * nodeIndex = [self _fetchNodeIndex];
-        NSMutableArray<NSNumber *> * nodeIndexCopy = [nodeIndex mutableCopy];
-        [nodeIndexCopy removeObject:nodeID];
-        if (nodeIndex.count != nodeIndexCopy.count) {
+
+        std::lock_guard lock(self->_nodeArrayLock);
+        auto oldCount = self->_nodesWithAttributeInfo.count;
+        [self->_nodesWithAttributeInfo removeObject:nodeID];
+        if (self->_nodesWithAttributeInfo.count != oldCount) {
             BOOL success;
-            if (nodeIndexCopy.count) {
-                success = [self _storeNodeIndex:nodeIndexCopy];
+            if (self->_nodesWithAttributeInfo.count) {
+                success = [self _storeNodeIndex:self->_nodesWithAttributeInfo];
             } else {
                 success = [self _deleteNodeIndex];
             }
             if (!success) {
-                MTR_LOG_ERROR("Store failed in clearStoredAttributesForNodeID for nodeIndex (%lu)", static_cast<unsigned long>(nodeIndexCopy.count));
+                MTR_LOG_ERROR("Store failed in clearStoredAttributesForNodeID for nodeIndex (%lu)", static_cast<unsigned long>(self->_nodesWithAttributeInfo.count));
             }
         }
     });
@@ -804,12 +845,12 @@ - (void)clearAllStoredClusterData
 {
     dispatch_async(_storageDelegateQueue, ^{
         // Fetch node index
-        NSArray<NSNumber *> * nodeIndex = [self _fetchNodeIndex];
-
-        for (NSNumber * nodeID in nodeIndex) {
+        std::lock_guard lock(self->_nodeArrayLock);
+        for (NSNumber * nodeID in self->_nodesWithAttributeInfo) {
             [self _clearStoredClusterDataForNodeID:nodeID];
         }
 
+        [self->_nodesWithAttributeInfo removeAllObjects];
         BOOL success = [self _deleteNodeIndex];
         if (!success) {
             MTR_LOG_ERROR("Delete failed for nodeIndex");
@@ -826,14 +867,13 @@ - (void)clearAllStoredClusterData
 
     __block NSMutableDictionary<MTRClusterPath *, MTRDeviceClusterData *> * clusterDataToReturn = nil;
     dispatch_sync(_storageDelegateQueue, ^{
-        // Fetch node index
-        NSArray<NSNumber *> * nodeIndex = [self _fetchNodeIndex];
+        std::lock_guard lock(self->_nodeArrayLock);
 
 #if ATTRIBUTE_CACHE_VERBOSE_LOGGING
-        MTR_LOG("Fetch got %lu values for nodeIndex", static_cast<unsigned long>(nodeIndex.count));
+        MTR_LOG("Fetch got %lu values for nodeIndex", static_cast<unsigned long>(self->_nodesWithAttributeInfo.count));
 #endif
 
-        if (![nodeIndex containsObject:nodeID]) {
+        if (![self->_nodesWithAttributeInfo containsObject:nodeID]) {
             // Sanity check and delete if nodeID exists in index
             NSArray<NSNumber *> * endpointIndex = [self _fetchEndpointIndexForNodeID:nodeID];
             if (endpointIndex) {
@@ -1086,14 +1126,13 @@ - (void)storeClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *>
         }
 
         // Check if node index needs updating / creation
-        NSArray<NSNumber *> * nodeIndex = [self _fetchNodeIndex];
         NSArray<NSNumber *> * nodeIndexToStore = nil;
-        if (!nodeIndex) {
-            // Ensure node index exists
-            MTR_LOG("No entry found for for nodeIndex - creating for node 0x%016llX", nodeID.unsignedLongLongValue);
-            nodeIndexToStore = [NSArray arrayWithObject:nodeID];
-        } else if (![nodeIndex containsObject:nodeID]) {
-            nodeIndexToStore = [nodeIndex arrayByAddingObject:nodeID];
+        {
+            std::lock_guard lock(self->_nodeArrayLock);
+            if (![self->_nodesWithAttributeInfo containsObject:nodeID]) {
+                [self->_nodesWithAttributeInfo addObject:nodeID];
+                nodeIndexToStore = [self->_nodesWithAttributeInfo copy];
+            }
         }
 
         if (nodeIndexToStore) {
@@ -1206,6 +1245,27 @@ - (void)synchronouslyPerformBlock:(void (^_Nullable)(void))block
     });
 }
 
+- (NSArray<NSNumber *> *)nodesWithStoredData
+{
+    // We have three types of stored data:
+    //
+    // 1) Attribute data
+    // 2) Session resumption data
+    // 3) Device data.
+    //
+    // Items 1 and 2 come with node indices.  Item 3 does not, but in practice
+    // we should have device data if and only if we have attribute data for that
+    // node ID, barring odd error conditions.
+    //
+    // TODO: Consider changing how we store device data so we can easily recover
+    // the relevant set of node IDs.
+    NSMutableSet<NSNumber *> * nodeSet = [NSMutableSet set];
+    std::lock_guard lock(_nodeArrayLock);
+    [nodeSet addObjectsFromArray:_nodesWithResumptionInfo];
+    [nodeSet addObjectsFromArray:_nodesWithAttributeInfo];
+    return [nodeSet allObjects];
+}
+
 @end
 
 @implementation MTRCASESessionResumptionInfo

src/darwin/Framework/CHIP/MTRDeviceController_Concrete.mm

@@ -1720,6 +1720,16 @@ - (nullable NSNumber *)neededReadPrivilegeForClusterID:(NSNumber *)clusterID att
     return nil;
 }
 
+- (NSArray<NSNumber *> *)nodesWithStoredData
+{
+    if (!self.controllerDataStore) {
+        // We have nothing stored, if we have no way to store.
+        return @[];
+    }
+
+    return [self.controllerDataStore nodesWithStoredData];
+}
+
 @end
 
 /**

src/darwin/Framework/CHIP/MTRDeviceController_XPC.mm

@@ -64,6 +64,11 @@ @implementation MTRDeviceController_XPC
                                                : (NSNumber *) nodeID, deleteNodeID
                                                : (NSNumber *) nodeID)
 
+MTR_DEVICECONTROLLER_SIMPLE_REMOTE_XPC_GETTER(nodesWithStoredData,
+    NSArray<NSNumber *> *,
+    @[], // Default return value
+    getNodesWithStoredDataWithReply)
+
 - (void)_updateRegistrationInfo
 {
     NSMutableDictionary * registrationInfo = [NSMutableDictionary dictionary];

src/darwin/Framework/CHIP/XPC Protocol/MTRXPCServerProtocol.h

@@ -77,6 +77,7 @@ MTR_AVAILABLE(ios(18.3), macos(15.3), watchos(11.3), tvos(18.3))
 - (oneway void)deviceController:(NSUUID *)controller unregisterNodeID:(NSNumber *)nodeID;
 - (oneway void)deviceController:(NSUUID *)controller updateControllerConfiguration:(NSDictionary *)controllerState MTR_AVAILABLE(ios(18.3), macos(15.3), watchos(11.3), tvos(18.3));
 
+- (oneway void)deviceController:(NSUUID *)controller getNodesWithStoredDataWithReply:(void (^)(NSArray<NSNumber *> *))reply MTR_AVAILABLE(ios(18.4), macos(15.4), watchos(11.4), tvos(18.4));
 @end
 
 MTR_AVAILABLE(ios(18.3), macos(15.3), watchos(11.3), tvos(18.3))

src/darwin/Framework/CHIPTests/MTRPerControllerStorageTests.m

@@ -1521,6 +1521,8 @@ - (void)test008_TestDataStoreDirect
     dispatch_sync(_storageQueue, ^{
         [storageDelegate controller:controller storeValues:testBulkValues securityLevel:MTRStorageSecurityLevelSecure sharingType:MTRStorageSharingTypeNotShared];
     });
+    // Since we messed with the node index, tell the data store to re-sync it's cache.
+    [controller.controllerDataStore unitTestRereadNodeIndex];
     // Verify that the store resulted in the correct values
     NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> * dataStoreClusterData = [controller.controllerDataStore getStoredClusterDataForNodeID:@(3001)];
     XCTAssertEqualObjects(dataStoreClusterData, bulkTestClusterDataDictionary);
@@ -1772,13 +1774,20 @@ - (void)test010_TestDataStoreMTRDeviceWithBulkReadWrite
     XCTAssertNotNil([dataStore findResumptionInfoByNodeID:deviceID]);
     XCTAssertNotNil([dataStore getStoredDeviceDataForNodeID:deviceID]);
     XCTAssertNotNil([dataStore getStoredClusterDataForNodeID:deviceID]);
+    __auto_type * nodesWithStoredData = [controller nodesWithStoredData];
+    XCTAssertTrue([nodesWithStoredData containsObject:deviceID]);
+    XCTAssertEqualObjects(nodesWithStoredData, [dataStore nodesWithStoredData]);
+    XCTAssertEqualObjects(nodesWithStoredData, deviceAttributeCounts.allKeys);
 
     [controller forgetDeviceWithNodeID:deviceID];
     deviceAttributeCounts = [controller unitTestGetDeviceAttributeCounts];
     XCTAssertNil(deviceAttributeCounts[deviceID]);
     XCTAssertNil([dataStore findResumptionInfoByNodeID:deviceID]);
     XCTAssertNil([dataStore getStoredDeviceDataForNodeID:deviceID]);
     XCTAssertNil([dataStore getStoredClusterDataForNodeID:deviceID]);
+    nodesWithStoredData = [controller nodesWithStoredData];
+    XCTAssertFalse([nodesWithStoredData containsObject:deviceID]);
+    XCTAssertEqualObjects(nodesWithStoredData, [dataStore nodesWithStoredData]);
 
     [controller shutdown];
     XCTAssertFalse([controller isRunning]);

src/darwin/Framework/CHIPTests/TestHelpers/MTRTestDeclarations.h

@@ -36,13 +36,15 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)clearStoredClusterDataForNodeID:(NSNumber *)nodeID;
 - (void)clearAllStoredClusterData;
 - (void)unitTestPruneEmptyStoredClusterDataBranches;
+- (void)unitTestRereadNodeIndex;
 - (NSString *)_endpointIndexKeyForNodeID:(NSNumber *)nodeID;
 - (NSString *)_clusterIndexKeyForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID;
 - (NSString *)_clusterDataKeyForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID clusterID:(NSNumber *)clusterID;
 - (nullable NSArray<NSNumber *> *)_fetchEndpointIndexForNodeID:(NSNumber *)nodeID;
 - (nullable NSArray<NSNumber *> *)_fetchClusterIndexForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID;
 - (nullable MTRDeviceClusterData *)_fetchClusterDataForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID clusterID:(NSNumber *)clusterID;
 - (nullable NSDictionary<NSString *, id> *)getStoredDeviceDataForNodeID:(NSNumber *)nodeID;
+- (NSArray<NSNumber *> *)nodesWithStoredData;
 @end
 
 // Declare internal methods for testing