From 7109708f07c2480224997216618de5e1a36099f7 Mon Sep 17 00:00:00 2001
From: Jeff Tung <100387939+jtung-apple@users.noreply.github.com>
Date: Wed, 24 May 2023 13:17:21 -0700
Subject: [PATCH] [Darwin] MTRAsyncCallbackWorkQueue tsan fix and API
 strengthening (#26800)

---
 .../CHIP/MTRAsyncCallbackWorkQueue.h          |  3 +
 .../CHIP/MTRAsyncCallbackWorkQueue.mm         | 75 ++++++++++++++++---
 .../CHIP/MTRAsyncCallbackWorkQueue_Internal.h |  3 -
 3 files changed, 68 insertions(+), 13 deletions(-)

diff --git a/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.h b/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.h
index bad4ded93833b0..acb6da8053327c 100644
--- a/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.h
+++ b/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.h
@@ -54,6 +54,7 @@ typedef void (^MTRAsyncCallbackReadyHandler)(id context, NSUInteger retryCount);
 - (void)invalidate;
 
 // Work items may be enqueued from any queue or thread
+// Note: Once a work item is enqueued, its handlers cannot be modified
 - (void)enqueueWorkItem:(MTRAsyncCallbackQueueWorkItem *)item;
 
 // TODO: Add a "set concurrency width" method to allow for more than 1 work item at a time
@@ -71,10 +72,12 @@ typedef void (^MTRAsyncCallbackReadyHandler)(id context, NSUInteger retryCount);
 
 // Called by the creater of the work item when async work is done and should
 // be removed from the queue. The work queue will run the next work item.
+// Note: This must only be called from within the readyHandler
 - (void)endWork;
 
 // Called by the creater of the work item when async work should be retried.
 // The work queue will call this workItem's readyHandler again.
+// Note: This must only be called from within the readyHandler
 - (void)retryWork;
 @end
 
diff --git a/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.mm b/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.mm
index e97ee154022cf6..2b31d07f004ec9 100644
--- a/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.mm
+++ b/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue.mm
@@ -41,10 +41,13 @@ - (void)retryWork:(MTRAsyncCallbackQueueWorkItem *)workItem;
 @end
 
 @interface MTRAsyncCallbackQueueWorkItem ()
+@property (nonatomic, readonly) os_unfair_lock lock;
 @property (nonatomic, strong, readonly) dispatch_queue_t queue;
 @property (nonatomic, readwrite) NSUInteger retryCount;
 @property (nonatomic, strong) MTRAsyncCallbackWorkQueue * workQueue;
+@property (nonatomic, readonly) BOOL enqueued;
 // Called by the queue
+- (void)markedEnqueued;
 - (void)callReadyHandlerWithContext:(id)context;
 - (void)cancel;
 @end
@@ -72,6 +75,13 @@ - (NSString *)description
 
 - (void)enqueueWorkItem:(MTRAsyncCallbackQueueWorkItem *)item
 {
+    if (item.enqueued) {
+        MTR_LOG_ERROR("MTRAsyncCallbackWorkQueue enqueueWorkItem: item cannot be enqueued twice");
+        return;
+    }
+
+    [item markedEnqueued];
+
     os_unfair_lock_lock(&_lock);
     item.workQueue = self;
     [self.items addObject:item];
@@ -163,12 +173,14 @@ @implementation MTRAsyncCallbackQueueWorkItem
 - (instancetype)initWithQueue:(dispatch_queue_t)queue
 {
     if (self = [super init]) {
+        _lock = OS_UNFAIR_LOCK_INIT;
         _queue = queue;
     }
     return self;
 }
 
-- (void)invalidate
+// assume lock is held
+- (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
@@ -181,6 +193,38 @@ - (void)invalidate
     _cancelHandler = nil;
 }
 
+- (void)invalidate
+{
+    os_unfair_lock_lock(&_lock);
+    [self _invalidate];
+    os_unfair_lock_unlock(&_lock);
+}
+
+- (void)markedEnqueued
+{
+    os_unfair_lock_lock(&_lock);
+    _enqueued = YES;
+    os_unfair_lock_unlock(&_lock);
+}
+
+- (void)setReadyHandler:(MTRAsyncCallbackReadyHandler)readyHandler
+{
+    os_unfair_lock_lock(&_lock);
+    if (!_enqueued) {
+        _readyHandler = readyHandler;
+    }
+    os_unfair_lock_unlock(&_lock);
+}
+
+- (void)setCancelHandler:(dispatch_block_t)cancelHandler
+{
+    os_unfair_lock_lock(&_lock);
+    if (!_enqueued) {
+        _cancelHandler = cancelHandler;
+    }
+    os_unfair_lock_unlock(&_lock);
+}
+
 - (void)endWork
 {
     [self.workQueue endWork:self];
@@ -196,12 +240,19 @@ - (void)retryWork
 - (void)callReadyHandlerWithContext:(id)context
 {
     dispatch_async(self.queue, ^{
-        if (self.readyHandler == nil) {
+        os_unfair_lock_lock(&self->_lock);
+        MTRAsyncCallbackReadyHandler readyHandler = self->_readyHandler;
+        NSUInteger retryCount = self->_retryCount;
+        if (readyHandler) {
+            self->_retryCount++;
+        }
+        os_unfair_lock_unlock(&self->_lock);
+
+        if (readyHandler == nil) {
             // Nothing to do here.
             [self endWork];
         } else {
-            self.readyHandler(context, self.retryCount);
-            self.retryCount++;
+            readyHandler(context, retryCount);
         }
     });
 }
@@ -209,11 +260,15 @@ - (void)callReadyHandlerWithContext:(id)context
 // Called by the work queue
 - (void)cancel
 {
-    dispatch_async(self.queue, ^{
-        if (self.cancelHandler != nil) {
-            self.cancelHandler();
-        }
-        [self invalidate];
-    });
+    os_unfair_lock_lock(&self->_lock);
+    dispatch_block_t cancelHandler = self->_cancelHandler;
+    [self _invalidate];
+    os_unfair_lock_unlock(&self->_lock);
+
+    if (cancelHandler) {
+        dispatch_async(self.queue, ^{
+            cancelHandler();
+        });
+    }
 }
 @end
diff --git a/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue_Internal.h b/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue_Internal.h
index 35eb765c3914ce..49e6343355a6d5 100644
--- a/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue_Internal.h
+++ b/src/darwin/Framework/CHIP/MTRAsyncCallbackWorkQueue_Internal.h
@@ -26,9 +26,6 @@ NS_ASSUME_NONNULL_BEGIN
 @interface MTRAsyncCallbackWorkQueue ()
 // The MTRDevice object is only held and passed back as a reference and is opaque to the queue
 - (instancetype)initWithContext:(id _Nullable)context queue:(dispatch_queue_t)queue;
-
-// Called by DeviceController at device clean up time
-- (void)invalidate;
 @end
 
 NS_ASSUME_NONNULL_END