Fix memory leaks in MTRAsyncCallbackQueueWorkItem. (project-chip#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.

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