From 23521629b8854103bc2e3ef2d170b9609bef726d Mon Sep 17 00:00:00 2001 From: Boris Zbarsky Date: Thu, 2 Feb 2023 12:44:29 -0500 Subject: [PATCH] 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. --- .../CHIP/MTRAsyncCallbackWorkQueue.mm | 28 +++++- .../CHIPTests/MTRAsyncCallbackQueueTests.m | 86 +++++++++++++++++++ 2 files changed, 111 insertions(+), 3 deletions(-) diff --git a/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.mm b/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.mm index edb2265dfc5db7..e97ee154022cf6 100644 --- a/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.mm +++ b/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,8 +196,13 @@ - (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++; + } }); } @@ -191,7 +210,10 @@ - (void)callReadyHandlerWithContext:(id)context - (void)cancel { dispatch_async(self.queue, ^{ - self.cancelHandler(); + if (self.cancelHandler != nil) { + self.cancelHandler(); + } + [self invalidate]; }); } @end diff --git a/src/darwin/Framework/CHIPTests/MTRAsyncCallbackQueueTests.m b/src/darwin/Framework/CHIPTests/MTRAsyncCallbackQueueTests.m index 7676d2051b9f5a..6612d8f2e3cebc 100644 --- a/src/darwin/Framework/CHIPTests/MTRAsyncCallbackQueueTests.m +++ b/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