From dc7149d7fc14630e36ce1615cb6152f3d7e7fd2c Mon Sep 17 00:00:00 2001
From: Jeff Tung <100387939+jtung-apple@users.noreply.github.com>
Date: Fri, 29 Mar 2024 12:42:22 -0700
Subject: [PATCH] [Darwin] MTRDevice attribute storage should write indexes
 only once each time

---
 .../CHIP/MTRDeviceControllerDataStore.mm      | 111 ++++++++++++------
 1 file changed, 74 insertions(+), 37 deletions(-)

diff --git a/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm b/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm
index 49a653e6e73bbe..4c7c43586aa29c 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm
+++ b/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm
@@ -983,64 +983,101 @@ - (void)storeClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *>
     dispatch_async(_storageDelegateQueue, ^{
         NSUInteger storeFailures = 0;
 
+        // 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];
+
+        // Write cluster specific data first
         for (MTRClusterPath * path in clusterData) {
             MTRDeviceClusterData * data = clusterData[path];
-
 #if ATTRIBUTE_CACHE_VERBOSE_LOGGING
             MTR_LOG_INFO("Attempt to store clusterData @ node 0x%016llX endpoint %u cluster 0x%08lX", nodeID.unsignedLongLongValue, path.endpoint.unsignedShortValue, path.cluster.unsignedLongValue);
 #endif
 
-            BOOL storeFailed = NO;
-            // Ensure node index exists
-            NSArray<NSNumber *> * nodeIndex = [self _fetchNodeIndex];
-            if (!nodeIndex) {
-                nodeIndex = [NSArray arrayWithObject:nodeID];
-                storeFailed = ![self _storeNodeIndex:nodeIndex];
-            } else if (![nodeIndex containsObject:nodeID]) {
-                storeFailed = ![self _storeNodeIndex:[nodeIndex arrayByAddingObject:nodeID]];
-            }
+            // 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 nodeIndex");
-                continue;
+                MTR_LOG_INFO("Store failed for clusterDAta @ node 0x%016llX endpoint %u cluster 0x%08lX", nodeID.unsignedLongLongValue, path.endpoint.unsignedShortValue, path.cluster.unsignedLongValue);
             }
 
-            // Ensure endpoint index exists
-            NSArray<NSNumber *> * endpointIndex = [self _fetchEndpointIndexForNodeID:nodeID];
-            if (!endpointIndex) {
-                endpointIndex = [NSArray arrayWithObject:path.endpoint];
-                storeFailed = ![self _storeEndpointIndex:endpointIndex forNodeID:nodeID];
-            } else if (![endpointIndex containsObject:path.endpoint]) {
-                storeFailed = ![self _storeEndpointIndex:[endpointIndex arrayByAddingObject:path.endpoint] forNodeID:nodeID];
+            // Note the cluster as modified for the endpoint
+            NSMutableSet<NSNumber *> * clustersInEndpoint = clustersModified[path.endpoint];
+            if (!clustersInEndpoint) {
+                clustersInEndpoint = [NSMutableSet set];
+                clustersModified[path.endpoint] = clustersInEndpoint;
             }
-            if (storeFailed) {
-                storeFailures++;
-                MTR_LOG_INFO("Store failed for endpointIndex @ node 0x%016llX", nodeID.unsignedLongLongValue);
-                continue;
+            [clustersInEndpoint addObject:path.cluster];
+        }
+
+        // Prepare to tally all endpoints touched
+        NSArray<NSNumber *> * endpointIndex = [self _fetchEndpointIndexForNodeID:nodeID];
+        BOOL endpointIndexModified = NO;
+        NSMutableArray<NSNumber *> * endpointIndexToStore;
+        if (endpointIndex) {
+            endpointIndexToStore = [endpointIndex mutableCopy];
+        } else {
+            endpointIndexToStore = [NSMutableArray array];
+            endpointIndexModified = YES;
+        }
+
+        for (NSNumber * endpointID in clustersModified) {
+            if (![endpointIndexToStore containsObject:endpointID]) {
+                [endpointIndexToStore addObject:endpointID];
+                endpointIndexModified = YES;
             }
 
-            // Ensure cluster index exists
-            NSArray<NSNumber *> * clusterIndex = [self _fetchClusterIndexForNodeID:nodeID endpointID:path.endpoint];
-            if (!clusterIndex) {
-                clusterIndex = [NSArray arrayWithObject:path.cluster];
-                storeFailed = ![self _storeClusterIndex:clusterIndex forNodeID:nodeID endpointID:path.endpoint];
-            } else if (![clusterIndex containsObject:path.cluster]) {
-                storeFailed = ![self _storeClusterIndex:[clusterIndex arrayByAddingObject:path.cluster] forNodeID:nodeID endpointID:path.endpoint];
+            // Get cluster index from storage and prepare to tally clusters touched
+            NSSet * newClusters = clustersModified[endpointID];
+            NSArray<NSNumber *> * clusterIndex = [self _fetchClusterIndexForNodeID:nodeID endpointID:endpointID];
+            BOOL clusterIndexModified = NO;
+            NSMutableArray<NSNumber *> * clusterIndexToStore;
+            if (clusterIndex) {
+                clusterIndexToStore = [clusterIndex mutableCopy];
+            } else {
+                clusterIndexToStore = [NSMutableArray array];
+                clusterIndexModified = YES;
             }
-            if (storeFailed) {
-                storeFailures++;
-                MTR_LOG_INFO("Store failed for clusterIndex @ node 0x%016llX endpoint %u", nodeID.unsignedLongLongValue, path.endpoint.unsignedShortValue);
-                continue;
+
+            for (NSNumber * clusterID in newClusters) {
+                if (![clusterIndexToStore containsObject:clusterID]) {
+                    [clusterIndexToStore addObject:clusterID];
+                    clusterIndexModified = YES;
+                }
             }
 
-            // Store cluster data
-            storeFailed = ![self _storeClusterData:data forNodeID:nodeID endpointID:path.endpoint clusterID:path.cluster];
+            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;
+                }
+            }
+        }
+
+        // Update endpoint index as needed
+        if (endpointIndexModified) {
+            BOOL storeFailed = ![self _storeEndpointIndex:endpointIndexToStore forNodeID:nodeID];
             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);
+                MTR_LOG_INFO("Store failed for endpointIndex @ node 0x%016llX", nodeID.unsignedLongLongValue);
             }
         }
 
+        // Ensure node index exists
+        NSArray<NSNumber *> * nodeIndex = [self _fetchNodeIndex];
+        BOOL storeFailed = NO;
+        if (!nodeIndex) {
+            nodeIndex = [NSArray arrayWithObject:nodeID];
+            storeFailed = ![self _storeNodeIndex:nodeIndex];
+        } else if (![nodeIndex containsObject:nodeID]) {
+            storeFailed = ![self _storeNodeIndex:[nodeIndex arrayByAddingObject:nodeID]];
+        }
+        if (storeFailed) {
+            storeFailures++;
+            MTR_LOG_INFO("Store failed for nodeIndex");
+        }
+
         // In the rare event that store fails, allow all attribute store attempts to go through and prune empty branches at the end altogether.
         if (storeFailures) {
             [self _pruneEmptyStoredAttributesBranches];