[Darwin] Unstored attributes should be flushed to storage on shutdown

src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.h

@@ -93,6 +93,16 @@ typedef void (^MTRDeviceControllerDataStoreClusterDataHandler)(NSDictionary<NSNu
 - (nullable NSDictionary<NSString *, id> *)getStoredDeviceDataForNodeID:(NSNumber *)nodeID;
 - (void)storeDeviceData:(NSDictionary<NSString *, id> *)data forNodeID:(NSNumber *)nodeID;
 
+/**
+ * Mechanism for and API client to perform a block after previous async operations (writes) on the storage queue have executed.
+ *
+ * This should be used only when something really needs to wait for the asynchronous writes
+ * to complete and can't proceed until they have.
+ *
+ * If no block is passed in, then the method returns after having synchronously flushed the queue.
+ */
+- (void)synchronouslyPerformBlock:(void (^_Nullable)(void))block;
+
 @end
 
 NS_ASSUME_NONNULL_END

src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm

@@ -1169,6 +1169,15 @@ - (void)storeDeviceData:(NSDictionary<NSString *, id> *)data forNodeID:(NSNumber
     });
 }
 
+- (void)synchronouslyPerformBlock:(void (^_Nullable)(void))block
+{
+    dispatch_sync(_storageDelegateQueue, ^{
+        if (block) {
+            block();
+        }
+    });
+}
+
 @end
 
 @implementation MTRCASESessionResumptionInfo

src/darwin/Framework/CHIP/MTRDeviceController_Concrete.mm

@@ -449,6 +449,13 @@ - (void)cleanupAfterStartup
     for (MTRDevice * device in devices) {
         [device invalidate];
     }
+
+    // Since MTRDevice invalidate may issue asynchronous writes to storage, perform a
+    // block synchronously on the storage delegate queue to flush those operations.
+    [self.controllerDataStore synchronouslyPerformBlock:^{
+        MTR_LOG("%@ Finished flushing data store", self);
+    }];
+
     [self stopBrowseForCommissionables];
 
     [_factory controllerShuttingDown:self];

src/darwin/Framework/CHIP/MTRDevice_Concrete.mm

@@ -850,6 +850,9 @@ - (void)invalidate
 
     os_unfair_lock_lock(&self->_lock);
 
+    // Flush unstored attributes if any
+    [self _persistClusterData];
+
     _state = MTRDeviceStateUnknown;
 
     // Make sure we don't try to resubscribe if we have a pending resubscribe

src/darwin/Framework/CHIPTests/MTRPerControllerStorageTests.m

@@ -1692,6 +1692,7 @@ - (void)doDataStoreMTRDeviceTestWithStorageDelegate:(id<MTRDeviceControllerStora
     [self waitForExpectations:@[ newDeviceSubscriptionExpectation ] timeout:60];
     if (!disableStorageBehaviorOptimization) {
         [self waitForExpectations:@[ newDeviceGotClusterDataPersisted ] timeout:60];
+        newDelegate.onClusterDataPersisted = nil;
     }
     newDelegate.onReportEnd = nil;
 
@@ -1953,6 +1954,128 @@ - (void)test013_suspendDevices
     [operationalBrowser shutdown];
 }
 
+- (void)test014_TestDataStoreMTRDeviceInvalidateFlush
+{
+    __auto_type * factory = [MTRDeviceControllerFactory sharedInstance];
+    XCTAssertNotNil(factory);
+
+    __auto_type queue = dispatch_get_main_queue();
+
+    __auto_type * rootKeys = [[MTRTestKeys alloc] init];
+    XCTAssertNotNil(rootKeys);
+
+    __auto_type * operationalKeys = [[MTRTestKeys alloc] init];
+    XCTAssertNotNil(operationalKeys);
+
+    NSNumber * nodeID = @(123);
+    NSNumber * fabricID = @(456);
+
+    NSError * error;
+
+    __auto_type * storageDelegate = [[MTRTestPerControllerStorageWithBulkReadWrite alloc] initWithControllerID:[NSUUID UUID]];
+
+    MTRPerControllerStorageTestsCertificateIssuer * certificateIssuer;
+    MTRDeviceStorageBehaviorConfiguration * storageBehaviorConfiguration = [MTRDeviceStorageBehaviorConfiguration configurationWithDefaultStorageBehavior];
+    MTRDeviceController * controller = [self startControllerWithRootKeys:rootKeys
+                                                         operationalKeys:operationalKeys
+                                                                fabricID:fabricID
+                                                                  nodeID:nodeID
+                                                                 storage:storageDelegate
+                                                                   error:&error
+                                                       certificateIssuer:&certificateIssuer
+                                            storageBehaviorConfiguration:storageBehaviorConfiguration];
+    XCTAssertNil(error);
+    XCTAssertNotNil(controller);
+    XCTAssertTrue([controller isRunning]);
+
+    XCTAssertEqualObjects(controller.controllerNodeID, nodeID);
+
+    // Now commission the device, to test that that works.
+    NSNumber * deviceID = @(17);
+    certificateIssuer.nextNodeID = deviceID;
+    [self commissionWithController:controller newNodeID:deviceID];
+
+    // We should have established CASE using our operational key.
+    XCTAssertEqual(operationalKeys.signatureCount, 1);
+
+    __auto_type * device = [MTRDevice deviceWithNodeID:deviceID controller:controller];
+    __auto_type * delegate = [[MTRDeviceTestDelegateWithSubscriptionSetupOverride alloc] init];
+
+    delegate.skipSetupSubscription = YES;
+
+    // Read the base storage key count (case session resumption etc.)
+    NSUInteger baseStorageKeyCount = storageDelegate.count;
+
+    [device setDelegate:delegate queue:queue];
+
+    NSArray<NSDictionary<NSString *, id> *> * attributeReport = @[ @{
+        MTRAttributePathKey : [MTRAttributePath attributePathWithEndpointID:@(0) clusterID:@(1) attributeID:@(1)],
+        MTRDataKey : @ {
+            MTRDataVersionKey : @(1),
+            MTRTypeKey : MTRUnsignedIntegerValueType,
+            MTRValueKey : @(1),
+        }
+    } ];
+
+    // Inject first report as priming report, which gets persisted immediately
+    [device unitTestInjectAttributeReport:attributeReport fromSubscription:YES];
+
+    // No additional entries immediately after injected report
+    XCTAssertEqual(storageDelegate.count, baseStorageKeyCount);
+
+    sleep(1);
+
+    // Verify priming report persisted before hitting storage delay
+    XCTAssertGreaterThan(storageDelegate.count, baseStorageKeyCount);
+    // Now set the base count to the after-priming number
+    baseStorageKeyCount = storageDelegate.count;
+
+    NSArray<NSDictionary<NSString *, id> *> * attributeReport2 = @[ @{
+        MTRAttributePathKey : [MTRAttributePath attributePathWithEndpointID:@(0) clusterID:@(2) attributeID:@(2)],
+        MTRDataKey : @ {
+            MTRDataVersionKey : @(2),
+            MTRTypeKey : MTRUnsignedIntegerValueType,
+            MTRValueKey : @(2),
+        }
+    } ];
+
+    // Inject second report with different cluster
+    [device unitTestInjectAttributeReport:attributeReport2 fromSubscription:YES];
+
+    sleep(1);
+
+    // No additional entries a second after report - under storage delay
+    XCTAssertEqual(storageDelegate.count, baseStorageKeyCount);
+
+    // Immediately shut down controller and force flush to storage
+    [controller shutdown];
+    XCTAssertFalse([controller isRunning]);
+
+    // Make sure there are more than base count entries
+    XCTAssertGreaterThan(storageDelegate.count, baseStorageKeyCount);
+
+    // Now restart controller to decommission the device
+    controller = [self startControllerWithRootKeys:rootKeys
+                                   operationalKeys:operationalKeys
+                                          fabricID:fabricID
+                                            nodeID:nodeID
+                                           storage:storageDelegate
+                                             error:&error
+                                 certificateIssuer:&certificateIssuer
+                      storageBehaviorConfiguration:storageBehaviorConfiguration];
+    XCTAssertNil(error);
+    XCTAssertNotNil(controller);
+    XCTAssertTrue([controller isRunning]);
+
+    XCTAssertEqualObjects(controller.controllerNodeID, nodeID);
+
+    // Reset our commissionee.
+    __auto_type * baseDevice = [MTRBaseDevice deviceWithNodeID:deviceID controller:controller];
+    ResetCommissionee(baseDevice, queue, self, kTimeoutInSeconds);
+
+    [controller shutdown];
+}
+
 // TODO: This might want to go in a separate test file, with some shared setup
 // across multiple tests, maybe.  Would need to factor out
 // startControllerWithRootKeys into a test helper.

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

@@ -38,6 +38,9 @@ NS_ASSUME_NONNULL_BEGIN
     removeValueForKey:(NSString *)key
         securityLevel:(MTRStorageSecurityLevel)securityLevel
           sharingType:(MTRStorageSharingType)sharingType;
+
+// For testing - direct access to the current count of keys in storage
+@property (nonatomic, readonly) NSUInteger count;
 @end
 
 @interface MTRTestPerControllerStorageWithBulkReadWrite : MTRTestPerControllerStorage

src/darwin/Framework/CHIPTests/TestHelpers/MTRTestPerControllerStorage.m

@@ -40,19 +40,21 @@ - (instancetype)initWithControllerID:(NSUUID *)controllerID
                             securityLevel:(MTRStorageSecurityLevel)securityLevel
                               sharingType:(MTRStorageSharingType)sharingType
 {
-    XCTAssertEqualObjects(_controllerID, controller.uniqueIdentifier);
+    @synchronized(self) {
+        XCTAssertEqualObjects(_controllerID, controller.uniqueIdentifier);
 
-    __auto_type * data = self.storage[key];
-    if (data == nil) {
-        return data;
-    }
+        __auto_type * data = self.storage[key];
+        if (data == nil) {
+            return data;
+        }
 
-    NSError * error;
-    id value = [NSKeyedUnarchiver unarchivedObjectOfClasses:MTRDeviceControllerStorageClasses() fromData:data error:&error];
-    XCTAssertNil(error);
-    XCTAssertNotNil(data);
+        NSError * error;
+        id value = [NSKeyedUnarchiver unarchivedObjectOfClasses:MTRDeviceControllerStorageClasses() fromData:data error:&error];
+        XCTAssertNil(error);
+        XCTAssertNotNil(data);
 
-    return value;
+        return value;
+    }
 }
 
 - (BOOL)controller:(MTRDeviceController *)controller
@@ -61,25 +63,36 @@ - (BOOL)controller:(MTRDeviceController *)controller
      securityLevel:(MTRStorageSecurityLevel)securityLevel
        sharingType:(MTRStorageSharingType)sharingType
 {
-    XCTAssertEqualObjects(_controllerID, controller.uniqueIdentifier);
+    @synchronized(self) {
+        XCTAssertEqualObjects(_controllerID, controller.uniqueIdentifier);
 
-    NSError * error;
-    NSData * data = [NSKeyedArchiver archivedDataWithRootObject:value requiringSecureCoding:YES error:&error];
-    XCTAssertNil(error);
-    XCTAssertNotNil(data);
+        NSError * error;
+        NSData * data = [NSKeyedArchiver archivedDataWithRootObject:value requiringSecureCoding:YES error:&error];
+        XCTAssertNil(error);
+        XCTAssertNotNil(data);
 
-    self.storage[key] = data;
-    return YES;
+        self.storage[key] = data;
+        return YES;
+    }
 }
 
 - (BOOL)controller:(MTRDeviceController *)controller
     removeValueForKey:(NSString *)key
         securityLevel:(MTRStorageSecurityLevel)securityLevel
           sharingType:(MTRStorageSharingType)sharingType
 {
-    XCTAssertEqualObjects(_controllerID, controller.uniqueIdentifier);
-    self.storage[key] = nil;
-    return YES;
+    @synchronized(self) {
+        XCTAssertEqualObjects(_controllerID, controller.uniqueIdentifier);
+        self.storage[key] = nil;
+        return YES;
+    }
+}
+
+- (NSUInteger)count
+{
+    @synchronized(self) {
+        return self.storage.count;
+    }
 }
 
 @end
@@ -88,29 +101,33 @@ @implementation MTRTestPerControllerStorageWithBulkReadWrite
 
 - (NSDictionary<NSString *, id<NSSecureCoding>> *)valuesForController:(MTRDeviceController *)controller securityLevel:(MTRStorageSecurityLevel)securityLevel sharingType:(MTRStorageSharingType)sharingType
 {
-    XCTAssertEqualObjects(self.controllerID, controller.uniqueIdentifier);
+    @synchronized(self) {
+        XCTAssertEqualObjects(self.controllerID, controller.uniqueIdentifier);
 
-    if (!self.storage.count) {
-        return nil;
-    }
+        if (!self.storage.count) {
+            return nil;
+        }
 
-    NSMutableDictionary * valuesToReturn = [NSMutableDictionary dictionary];
-    for (NSString * key in self.storage) {
-        valuesToReturn[key] = [self controller:controller valueForKey:key securityLevel:securityLevel sharingType:sharingType];
-    }
+        NSMutableDictionary * valuesToReturn = [NSMutableDictionary dictionary];
+        for (NSString * key in self.storage) {
+            valuesToReturn[key] = [self controller:controller valueForKey:key securityLevel:securityLevel sharingType:sharingType];
+        }
 
-    return valuesToReturn;
+        return valuesToReturn;
+    }
 }
 
 - (BOOL)controller:(MTRDeviceController *)controller storeValues:(NSDictionary<NSString *, id<NSSecureCoding>> *)values securityLevel:(MTRStorageSecurityLevel)securityLevel sharingType:(MTRStorageSharingType)sharingType
 {
-    XCTAssertEqualObjects(self.controllerID, controller.uniqueIdentifier);
+    @synchronized(self) {
+        XCTAssertEqualObjects(self.controllerID, controller.uniqueIdentifier);
 
-    for (NSString * key in values) {
-        [self controller:controller storeValue:values[key] forKey:key securityLevel:securityLevel sharingType:sharingType];
-    }
+        for (NSString * key in values) {
+            [self controller:controller storeValue:values[key] forKey:key securityLevel:securityLevel sharingType:sharingType];
+        }
 
-    return YES;
+        return YES;
+    }
 }
 
 @end