Fix memory leaks in MTRAsyncCallbackQueueWorkItem. (#24817)

The item API pretty much requires a readyHandler that holds a reference to the
item.  Instead of assuming all consumers use __weak correctly to hold that
reference, manually break cycles in MTRAsyncCallbackQueueWorkItem.

Author: bzbarsky-apple

src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.mm

@@ -168,9 +168,23 @@ - (instancetype)initWithQueue:(dispatch_queue_t)queue
     return self;
 }
 
+- (void)invalidate
+{
+    // Make sure we don't leak via handlers that close over us, as ours must.
+    // This is a bit odd, since these are supposed to be non-nullable
+    // properties, but it's the best we can do given our API surface, unless we
+    // assume that all consumers consistently use __weak refs to us inside their
+    // handlers.
+    //
+    // Setting the attributes to nil will not compile; set the ivars directly.
+    _readyHandler = nil;
+    _cancelHandler = nil;
+}
+
 - (void)endWork
 {
     [self.workQueue endWork:self];
+    [self invalidate];
 }
 
 - (void)retryWork
@@ -182,16 +196,24 @@ - (void)retryWork
 - (void)callReadyHandlerWithContext:(id)context
 {
     dispatch_async(self.queue, ^{
-        self.readyHandler(context, self.retryCount);
-        self.retryCount++;
+        if (self.readyHandler == nil) {
+            // Nothing to do here.
+            [self endWork];
+        } else {
+            self.readyHandler(context, self.retryCount);
+            self.retryCount++;
+        }
     });
 }
 
 // Called by the work queue
 - (void)cancel
 {
     dispatch_async(self.queue, ^{
-        self.cancelHandler();
+        if (self.cancelHandler != nil) {
+            self.cancelHandler();
+        }
+        [self invalidate];
     });
 }
 @end

src/darwin/Framework/CHIPTests/MTRAsyncCallbackQueueTests.m

@@ -46,6 +46,12 @@ - (void)testRunItem
     };
     [workQueue enqueueWorkItem:workItem1];
 
+    // Check for leaks.
+    MTRAsyncCallbackQueueWorkItem * __weak weakItem = workItem1;
+    [self addTeardownBlock:^() {
+        XCTAssertNil(weakItem);
+    }];
+
     [self waitForExpectationsWithTimeout:5 handler:nil];
 
     // see that it only ran once
@@ -170,4 +176,84 @@ - (void)testRunItemsAfterDrain
     [self waitForExpectationsWithTimeout:5 handler:nil];
 }
 
+- (void)testRunItemNoHandlers
+{
+    XCTestExpectation * expectation = [self expectationWithDescription:@"Work item called"];
+
+    MTRAsyncCallbackWorkQueue * workQueue = [[MTRAsyncCallbackWorkQueue alloc] initWithContext:nil queue:dispatch_get_main_queue()];
+
+    MTRAsyncCallbackQueueWorkItem * workItem1 =
+        [[MTRAsyncCallbackQueueWorkItem alloc] initWithQueue:dispatch_get_global_queue(QOS_CLASS_DEFAULT, 0)];
+    MTRAsyncCallbackQueueWorkItem * workItem2 =
+        [[MTRAsyncCallbackQueueWorkItem alloc] initWithQueue:dispatch_get_global_queue(QOS_CLASS_DEFAULT, 0)];
+
+    __block int counter = 0;
+    MTRAsyncCallbackReadyHandler readyHandler = ^(MTRDevice * _Nonnull device, NSUInteger retryCount) {
+        counter++;
+        [workItem2 endWork];
+        [expectation fulfill];
+    };
+    workItem2.readyHandler = readyHandler;
+    workItem2.cancelHandler = ^{
+    };
+
+    // Check that trying to run workItem1 does not crash.
+    [workQueue enqueueWorkItem:workItem1];
+    [workQueue enqueueWorkItem:workItem2];
+
+    [self waitForExpectationsWithTimeout:5 handler:nil];
+
+    // see that it only ran once
+    XCTAssertEqual(counter, 1);
+}
+
+- (void)testInvalidation
+{
+    XCTestExpectation * expectation = [self expectationWithDescription:@"Work item called"];
+    XCTestExpectation * cancelExpectation = [self expectationWithDescription:@"Work item canceled"];
+
+    MTRAsyncCallbackWorkQueue * workQueue = [[MTRAsyncCallbackWorkQueue alloc] initWithContext:nil queue:dispatch_get_main_queue()];
+
+    MTRAsyncCallbackQueueWorkItem * workItem1 =
+        [[MTRAsyncCallbackQueueWorkItem alloc] initWithQueue:dispatch_get_global_queue(QOS_CLASS_DEFAULT, 0)];
+    MTRAsyncCallbackReadyHandler readyHandler1 = ^(MTRDevice * _Nonnull device, NSUInteger retryCount) {
+        // Give the code enqueing the other items a chance to run, so they can
+        // actually get canceled.
+        sleep(1);
+        [workQueue invalidate];
+        [workItem1 endWork];
+        [expectation fulfill];
+    };
+    workItem1.readyHandler = readyHandler1;
+    // No cancel handler on purpose.
+    [workQueue enqueueWorkItem:workItem1];
+
+    MTRAsyncCallbackQueueWorkItem * workItem2 =
+        [[MTRAsyncCallbackQueueWorkItem alloc] initWithQueue:dispatch_get_global_queue(QOS_CLASS_DEFAULT, 0)];
+    MTRAsyncCallbackReadyHandler readyHandler2 = ^(MTRDevice * _Nonnull device, NSUInteger retryCount) {
+        // This should never get called.
+        XCTAssertFalse(YES);
+        [workItem2 endWork];
+    };
+    workItem2.readyHandler = readyHandler2;
+    // No cancel handler on purpose.
+    [workQueue enqueueWorkItem:workItem2];
+
+    MTRAsyncCallbackQueueWorkItem * workItem3 =
+        [[MTRAsyncCallbackQueueWorkItem alloc] initWithQueue:dispatch_get_global_queue(QOS_CLASS_DEFAULT, 0)];
+    MTRAsyncCallbackReadyHandler readyHandler3 = ^(MTRDevice * _Nonnull device, NSUInteger retryCount) {
+        // This should never get called.
+        XCTAssertFalse(YES);
+        [workItem3 endWork];
+    };
+    dispatch_block_t cancelHandler3 = ^() {
+        [cancelExpectation fulfill];
+    };
+    workItem3.readyHandler = readyHandler3;
+    workItem3.cancelHandler = cancelHandler3;
+    [workQueue enqueueWorkItem:workItem3];
+
+    [self waitForExpectations:@[ expectation, cancelExpectation ] timeout:5];
+}
+
 @end