[Darwin] MTRDeviceControllerStorageDelegate should support bulk store/fetch

src/darwin/Framework/CHIP/MTRDevice.mm

@@ -236,6 +236,20 @@ - (id)copyWithZone:(NSZone *)zone
     return [[MTRDeviceClusterData alloc] initWithDataVersion:_dataVersion attributes:_attributes];
 }
 
+- (BOOL)isEqualToClusterData:(MTRDeviceClusterData *)otherClusterData
+{
+    return [_dataVersion isEqual:otherClusterData.dataVersion] && [_attributes isEqual:otherClusterData.attributes];
+}
+
+- (BOOL)isEqual:(id)object
+{
+    if ([object class] != [self class]) {
+        return NO;
+    }
+
+    return [self isEqualToClusterData:object];
+}
+
 @end
 
 @interface MTRDevice ()
@@ -2241,6 +2255,14 @@ - (void)setAttributeValues:(NSArray<NSDictionary *> *)attributeValues reportChan
     [self _setAttributeValues:attributeValues reportChanges:reportChanges];
 }
 
+#ifdef DEBUG
+- (NSUInteger)unitTestAttributeCount
+{
+    std::lock_guard lock(_lock);
+    return _readCache.count;
+}
+#endif
+
 - (void)setClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)clusterData
 {
     MTR_LOG_INFO("%@ setClusterData count: %lu", self, static_cast<unsigned long>(clusterData.count));

src/darwin/Framework/CHIP/MTRDeviceController.mm

@@ -590,6 +590,20 @@ - (BOOL)startup:(MTRDeviceControllerStartupParamsInternal *)startupParams
         return NO;
     }
 
+    if (_controllerDataStore) {
+        // If the storage delegate supports the bulk read API, then a dictionary of nodeID => cluster data dictionary would be passed to the handler. Otherwise this would be a no-op, and stored attributes for MTRDevice objects will be loaded lazily in -deviceForNodeID:.
+        [_controllerDataStore fetchAttributeDataForAllDevices:^(NSDictionary<NSNumber *, NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *> * _Nonnull clusterDataByNode) {
+            MTR_LOG_INFO("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);
+            for (NSNumber * nodeID in clusterDataByNode) {
+                NSDictionary * clusterData = clusterDataByNode[nodeID];
+                MTRDevice * device = [self _setupDeviceForNodeID:nodeID prefetchedClusterData:clusterData];
+                MTR_LOG_INFO("Loaded %lu cluster data from storage for %@", static_cast<unsigned long>(clusterData.count), device);
+            }
+        }];
+    }
+
     return YES;
 }
 
@@ -933,20 +947,25 @@ - (MTRBaseDevice *)baseDeviceForNodeID:(NSNumber *)nodeID
     return [[MTRBaseDevice alloc] initWithNodeID:nodeID controller:self];
 }
 
-- (MTRDevice *)deviceForNodeID:(NSNumber *)nodeID
+// If prefetchedClusterData is not provided, load attributes individually from controller data store
+- (MTRDevice *)_setupDeviceForNodeID:(NSNumber *)nodeID prefetchedClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)prefetchedClusterData
 {
-    std::lock_guard lock(_deviceMapLock);
-    MTRDevice * deviceToReturn = _nodeIDToDeviceMap[nodeID];
-    if (!deviceToReturn) {
-        deviceToReturn = [[MTRDevice alloc] initWithNodeID:nodeID controller:self];
-        // If we're not running, don't add the device to our map.  That would
-        // create a cycle that nothing would break.  Just return the device,
-        // which will be in exactly the state it would be in if it were created
-        // while we were running and then we got shut down.
-        if ([self isRunning]) {
-            _nodeIDToDeviceMap[nodeID] = deviceToReturn;
-        }
+    os_unfair_lock_assert_owner(&_deviceMapLock);
+
+    MTRDevice * deviceToReturn = [[MTRDevice alloc] initWithNodeID:nodeID controller:self];
+    // If we're not running, don't add the device to our map.  That would
+    // create a cycle that nothing would break.  Just return the device,
+    // which will be in exactly the state it would be in if it were created
+    // while we were running and then we got shut down.
+    if ([self isRunning]) {
+        _nodeIDToDeviceMap[nodeID] = deviceToReturn;
+    }
 
+    if (prefetchedClusterData) {
+        if (prefetchedClusterData.count) {
+            [deviceToReturn setClusterData:prefetchedClusterData];
+        }
+    } else {
 #if !MTRDEVICE_ATTRIBUTE_CACHE_STORE_ATTRIBUTES_BY_CLUSTER
         // Load persisted attributes if they exist.
         NSArray * attributesFromCache = [_controllerDataStore getStoredAttributesForNodeID:nodeID];
@@ -966,6 +985,17 @@ - (MTRDevice *)deviceForNodeID:(NSNumber *)nodeID
     return deviceToReturn;
 }
 
+- (MTRDevice *)deviceForNodeID:(NSNumber *)nodeID
+{
+    std::lock_guard lock(_deviceMapLock);
+    MTRDevice * deviceToReturn = _nodeIDToDeviceMap[nodeID];
+    if (!deviceToReturn) {
+        deviceToReturn = [self _setupDeviceForNodeID:nodeID prefetchedClusterData:nil];
+    }
+
+    return deviceToReturn;
+}
+
 - (void)removeDevice:(MTRDevice *)device
 {
     std::lock_guard lock(_deviceMapLock);
@@ -979,6 +1009,18 @@ - (void)removeDevice:(MTRDevice *)device
     }
 }
 
+#ifdef DEBUG
+- (NSDictionary<NSNumber *, NSNumber *> *)unitTestGetDeviceAttributeCounts
+{
+    std::lock_guard lock(_deviceMapLock);
+    NSMutableDictionary<NSNumber *, NSNumber *> * deviceAttributeCounts = [NSMutableDictionary dictionary];
+    for (NSNumber * nodeID in _nodeIDToDeviceMap) {
+        deviceAttributeCounts[nodeID] = @([_nodeIDToDeviceMap[nodeID] unitTestAttributeCount]);
+    }
+    return deviceAttributeCounts;
+}
+#endif
+
 - (void)setDeviceControllerDelegate:(id<MTRDeviceControllerDelegate>)delegate queue:(dispatch_queue_t)queue
 {
     [self

src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.h

@@ -48,6 +48,14 @@ NS_ASSUME_NONNULL_BEGIN
                             storageDelegate:(id<MTRDeviceControllerStorageDelegate>)storageDelegate
                        storageDelegateQueue:(dispatch_queue_t)storageDelegateQueue;
 
+// clusterDataByNode a dictionary: nodeID => cluster data dictionary
+typedef void (^MTRDeviceControllerDataStoreClusterDataHandler)(NSDictionary<NSNumber *, NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *> * clusterDataByNode);
+
+/**
+ * Asks the data store to load cluster data for nodes in bulk. If the storageDelegate supports it, the handler will be called synchronously.
+ */
+- (void)fetchAttributeDataForAllDevices:(MTRDeviceControllerDataStoreClusterDataHandler)clusterDataHandler;
+
 /**
  * Resumption info APIs.
  */

src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm

@@ -151,6 +151,20 @@ - (nullable instancetype)initWithController:(MTRDeviceController *)controller
     return self;
 }
 
+- (void)fetchAttributeDataForAllDevices:(MTRDeviceControllerDataStoreClusterDataHandler)clusterDataHandler
+{
+    __block NSDictionary<NSNumber *, NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *> * clusterDataByNode = nil;
+    dispatch_sync(_storageDelegateQueue, ^{
+        if ([self->_storageDelegate respondsToSelector:@selector(valuesForController:securityLevel:sharingType:)]) {
+            NSDictionary<NSString *, id> * dataStoreSecureLocalValues = [self->_storageDelegate valuesForController:self->_controller securityLevel:MTRStorageSecurityLevelSecure sharingType:MTRStorageSharingTypeNotShared];
+
+            clusterDataByNode = [self _getClusterDataFromSecureLocalValues:dataStoreSecureLocalValues];
+        }
+    });
+
+    clusterDataHandler(clusterDataByNode);
+}
+
 - (nullable MTRCASESessionResumptionInfo *)findResumptionInfoByNodeID:(NSNumber *)nodeID
 {
     return [self _findResumptionInfoWithKey:ResumptionByNodeIDKey(nodeID)];
@@ -337,6 +351,14 @@ - (BOOL)_storeAttributeCacheValue:(id)value forKey:(NSString *)key
                             sharingType:MTRStorageSharingTypeNotShared];
 }
 
+- (BOOL)_bulkStoreAttributeCacheValues:(NSDictionary<NSString *, id<NSSecureCoding>> *)values
+{
+    return [_storageDelegate controller:_controller
+                            storeValues:values
+                          securityLevel:MTRStorageSecurityLevelSecure
+                            sharingType:MTRStorageSharingTypeNotShared];
+}
+
 - (BOOL)_removeAttributeCacheValueForKey:(NSString *)key
 {
     return [_storageDelegate controller:_controller
@@ -968,6 +990,76 @@ - (void)clearAllStoredAttributes
     return clusterDataToReturn;
 }
 
+// Utility for constructing dictionary of nodeID to cluster data from dictionary of storage keys
+- (nullable NSDictionary<NSNumber *, NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *> *)_getClusterDataFromSecureLocalValues:(NSDictionary<NSString *, id> *)secureLocalValues
+{
+    NSMutableDictionary<NSNumber *, NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *> * clusterDataByNodeToReturn = nil;
+
+    if (![secureLocalValues isKindOfClass:[NSDictionary class]]) {
+        return nil;
+    }
+
+    // Fetch node index
+    NSArray<NSNumber *> * nodeIndex = secureLocalValues[sAttributeCacheNodeIndexKey];
+
+    if (![nodeIndex isKindOfClass:[NSArray class]]) {
+        return nil;
+    }
+
+    for (NSNumber * nodeID in nodeIndex) {
+        if (![nodeID isKindOfClass:[NSNumber class]]) {
+            continue;
+        }
+
+        NSMutableDictionary<MTRClusterPath *, MTRDeviceClusterData *> * clusterDataForNode = nil;
+        NSArray<NSNumber *> * endpointIndex = secureLocalValues[[self _endpointIndexKeyForNodeID:nodeID]];
+
+        if (![endpointIndex isKindOfClass:[NSArray class]]) {
+            continue;
+        }
+
+        for (NSNumber * endpointID in endpointIndex) {
+            if (![endpointID isKindOfClass:[NSNumber class]]) {
+                continue;
+            }
+
+            NSArray<NSNumber *> * clusterIndex = secureLocalValues[[self _clusterIndexKeyForNodeID:nodeID endpointID:endpointID]];
+
+            if (![clusterIndex isKindOfClass:[NSArray class]]) {
+                continue;
+            }
+
+            for (NSNumber * clusterID in clusterIndex) {
+                if (![clusterID isKindOfClass:[NSNumber class]]) {
+                    continue;
+                }
+
+                MTRDeviceClusterData * clusterData = secureLocalValues[[self _clusterDataKeyForNodeID:nodeID endpointID:endpointID clusterID:clusterID]];
+                if (!clusterData) {
+                    continue;
+                }
+                if (![clusterData isKindOfClass:[MTRDeviceClusterData class]]) {
+                    continue;
+                }
+                MTRClusterPath * clusterPath = [MTRClusterPath clusterPathWithEndpointID:endpointID clusterID:clusterID];
+                if (!clusterDataForNode) {
+                    clusterDataForNode = [NSMutableDictionary dictionary];
+                }
+                clusterDataForNode[clusterPath] = clusterData;
+            }
+        }
+
+        if (clusterDataForNode.count) {
+            if (!clusterDataByNodeToReturn) {
+                clusterDataByNodeToReturn = [NSMutableDictionary dictionary];
+            }
+            clusterDataByNodeToReturn[nodeID] = clusterDataForNode;
+        }
+    }
+
+    return clusterDataByNodeToReturn;
+}
+
 - (void)storeClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)clusterData forNodeID:(NSNumber *)nodeID
 {
     if (!nodeID) {
@@ -983,6 +1075,11 @@ - (void)storeClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *>
     dispatch_async(_storageDelegateQueue, ^{
         NSUInteger storeFailures = 0;
 
+        NSMutableDictionary<NSString *, id<NSSecureCoding>> * bulkValuesToStore = nil;
+        if ([self->_storageDelegate respondsToSelector:@selector(controller:storeValues:securityLevel:sharingType:)]) {
+            bulkValuesToStore = [NSMutableDictionary dictionary];
+        }
+
         // A map of endpoint => list of clusters modified for that endpoint so cluster indexes can be updated later
         NSMutableDictionary<NSNumber *, NSMutableSet<NSNumber *> *> * clustersModified = [NSMutableDictionary dictionary];
 
@@ -993,11 +1090,15 @@ - (void)storeClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *>
             MTR_LOG_INFO("Attempt to store clusterData @ node 0x%016llX endpoint %u cluster 0x%08lX", nodeID.unsignedLongLongValue, path.endpoint.unsignedShortValue, path.cluster.unsignedLongValue);
 #endif
 
-            // Store cluster data
-            BOOL storeFailed = ![self _storeClusterData:data forNodeID:nodeID endpointID:path.endpoint clusterID:path.cluster];
-            if (storeFailed) {
-                storeFailures++;
-                MTR_LOG_INFO("Store failed for clusterDAta @ node 0x%016llX endpoint %u cluster 0x%08lX", nodeID.unsignedLongLongValue, path.endpoint.unsignedShortValue, path.cluster.unsignedLongValue);
+            if (bulkValuesToStore) {
+                bulkValuesToStore[[self _clusterDataKeyForNodeID:nodeID endpointID:path.endpoint clusterID:path.cluster]] = data;
+            } else {
+                // Store cluster data
+                BOOL storeFailed = ![self _storeClusterData:data forNodeID:nodeID endpointID:path.endpoint clusterID:path.cluster];
+                if (storeFailed) {
+                    storeFailures++;
+                    MTR_LOG_INFO("Store failed for clusterData @ node 0x%016llX endpoint %u cluster 0x%08lX", nodeID.unsignedLongLongValue, path.endpoint.unsignedShortValue, path.cluster.unsignedLongValue);
+                }
             }
 
             // Note the cluster as modified for the endpoint
@@ -1046,36 +1147,60 @@ - (void)storeClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *>
             }
 
             if (clusterIndexModified) {
-                BOOL storeFailed = ![self _storeClusterIndex:clusterIndexToStore forNodeID:nodeID endpointID:endpointID];
-                if (storeFailed) {
-                    storeFailures++;
-                    MTR_LOG_INFO("Store failed for clusterIndex @ node 0x%016llX endpoint %u", nodeID.unsignedLongLongValue, endpointID.unsignedShortValue);
-                    continue;
+                if (bulkValuesToStore) {
+                    bulkValuesToStore[[self _clusterIndexKeyForNodeID:nodeID endpointID:endpointID]] = clusterIndexToStore;
+                } else {
+                    BOOL storeFailed = ![self _storeClusterIndex:clusterIndexToStore forNodeID:nodeID endpointID:endpointID];
+                    if (storeFailed) {
+                        storeFailures++;
+                        MTR_LOG_INFO("Store failed for clusterIndex @ node 0x%016llX endpoint %u", nodeID.unsignedLongLongValue, endpointID.unsignedShortValue);
+                        continue;
+                    }
                 }
             }
         }
 
         // Update endpoint index as needed
         if (endpointIndexModified) {
-            BOOL storeFailed = ![self _storeEndpointIndex:endpointIndexToStore forNodeID:nodeID];
-            if (storeFailed) {
-                storeFailures++;
-                MTR_LOG_INFO("Store failed for endpointIndex @ node 0x%016llX", nodeID.unsignedLongLongValue);
+            if (bulkValuesToStore) {
+                bulkValuesToStore[[self _endpointIndexKeyForNodeID:nodeID]] = endpointIndexToStore;
+            } else {
+                BOOL storeFailed = ![self _storeEndpointIndex:endpointIndexToStore forNodeID:nodeID];
+                if (storeFailed) {
+                    storeFailures++;
+                    MTR_LOG_INFO("Store failed for endpointIndex @ node 0x%016llX", nodeID.unsignedLongLongValue);
+                }
             }
         }
 
-        // Ensure node index exists
+        // Check if node index needs updating / creation
         NSArray<NSNumber *> * nodeIndex = [self _fetchNodeIndex];
-        BOOL storeFailed = NO;
+        NSArray<NSNumber *> * nodeIndexToStore = nil;
         if (!nodeIndex) {
-            nodeIndex = [NSArray arrayWithObject:nodeID];
-            storeFailed = ![self _storeNodeIndex:nodeIndex];
+            // Ensure node index exists
+            nodeIndexToStore = [NSArray arrayWithObject:nodeID];
         } else if (![nodeIndex containsObject:nodeID]) {
-            storeFailed = ![self _storeNodeIndex:[nodeIndex arrayByAddingObject:nodeID]];
+            nodeIndexToStore = [nodeIndex arrayByAddingObject:nodeID];
         }
-        if (storeFailed) {
-            storeFailures++;
-            MTR_LOG_INFO("Store failed for nodeIndex");
+
+        if (nodeIndexToStore) {
+            if (bulkValuesToStore) {
+                bulkValuesToStore[sAttributeCacheNodeIndexKey] = nodeIndexToStore;
+            } else {
+                BOOL storeFailed = ![self _storeNodeIndex:nodeIndexToStore];
+                if (storeFailed) {
+                    storeFailures++;
+                    MTR_LOG_INFO("Store failed for nodeIndex");
+                }
+            }
+        }
+
+        if (bulkValuesToStore) {
+            BOOL storeFailed = ![self _bulkStoreAttributeCacheValues:bulkValuesToStore];
+            if (storeFailed) {
+                storeFailures++;
+                MTR_LOG_INFO("Store failed for bulk values count %lu", static_cast<unsigned long>(bulkValuesToStore.count));
+            }
         }
 
         // In the rare event that store fails, allow all attribute store attempts to go through and prune empty branches at the end altogether.