From 92257905e1c40b34736cdd03b431716cce436e2f Mon Sep 17 00:00:00 2001
From: Darren Mo <darrenmo@google.com>
Date: Fri, 24 Jun 2022 16:43:04 -0700
Subject: [PATCH] Make expectation recording thread-safe.

To handle a mocked invocation, we do the following:
1. Get the first stub that matches the invocation.
2. If the stub is an expectation, remove the stub from the arrays of stubs and expectations.
3. Pass the invocation to the stub.

The first two steps are currently not atomic, so there is a race condition. For example:
1. Add expectation 1 for method A.
2. Add expectation 2 for same method A.
3. [Thread 1] Call method A.
4. [Thread 2] Call method A.
5. [Thread 1] Get expectation 1 since it matches the invocation.
6. [Thread 2] Get expectation 1 since it matches the invocation.
7. [Thread 1] Remove expectation 1 from the arrays.
8. [Thread 2] Remove expectation 1 from the arrays.

In the above example, the invocations of the same method in Thread 1 and Thread 2 both match expectation 1; expectation 2 does not get matched.

The solution is to make the matching and the removal atomic. We can do so by wrapping the relevant code in `@synchronized(stubs)`. (Even though `stubForInvocation:` also performs `@synchronized(stubs)`, we can feel free to do the same in the caller without fear of deadlocks since `@synchronized` uses a recursive lock.)

Added a new unit test that was failing before the fix and is now passing after the fix.
---
 Source/OCMock/OCMockObject.m           | 56 ++++++++++++--------------
 Source/OCMockTests/OCMockObjectTests.m | 19 +++++++++
 2 files changed, 45 insertions(+), 30 deletions(-)

diff --git a/Source/OCMock/OCMockObject.m b/Source/OCMock/OCMockObject.m
index d78cba44..7380e1bd 100644
--- a/Source/OCMock/OCMockObject.m
+++ b/Source/OCMock/OCMockObject.m
@@ -408,44 +408,40 @@ - (BOOL)handleInvocation:(NSInvocation *)anInvocation
     [self assertInvocationsArrayIsPresent];
     [self addInvocation:anInvocation];
 
-    OCMInvocationStub *stub = [self stubForInvocation:anInvocation];
-    if(stub == nil)
-        return NO;
+    OCMInvocationStub *stub = nil;
+    @synchronized(stubs)
+    {
+        stub = [self stubForInvocation:anInvocation];
+        if(stub == nil)
+            return NO;
 
-    // Retain the stub in case it ends up being removed because we still need it at the end for handleInvocation:
-    [stub retain];
+        // Retain the stub in case it ends up being removed because we still need it at the end for handleInvocation:
+        [stub retain];
 
-    BOOL removeStub = NO;
-    @synchronized(expectations)
-    {
-        if([expectations containsObject:stub])
+        @synchronized(expectations)
         {
-            OCMInvocationExpectation *expectation = [self _nextExpectedInvocation];
-            if(expectationOrderMatters && (expectation != stub))
+            if([expectations containsObject:stub])
             {
-                [NSException raise:NSInternalInconsistencyException
-                            format:@"%@: unexpected method invoked: %@\n\texpected:\t%@",
-                            [self description], [stub description], [[expectations objectAtIndex:0] description]];
-            }
+                OCMInvocationExpectation *expectation = [self _nextExpectedInvocation];
+                if(expectationOrderMatters && (expectation != stub))
+                {
+                    [NSException raise:NSInternalInconsistencyException
+                                format:@"%@: unexpected method invoked: %@\n\texpected:\t%@",
+                                [self description], [stub description], [[expectations objectAtIndex:0] description]];
+                }
 
-            // We can't check isSatisfied yet, since the stub won't be satisfied until we call
-            // handleInvocation: since we'll still have the current expectation in the expectations array, which
-            // will cause an exception if expectationOrderMatters is YES and we're not ready for any future
-            // expected methods to be called yet
-            if(![(OCMInvocationExpectation *)stub isMatchAndReject])
-            {
-                [expectations removeObject:stub];
-                removeStub = YES;
+                // We can't check isSatisfied yet, since the stub won't be satisfied until we call
+                // handleInvocation: since we'll still have the current expectation in the expectations array, which
+                // will cause an exception if expectationOrderMatters is YES and we're not ready for any future
+                // expected methods to be called yet
+                if(![(OCMInvocationExpectation *)stub isMatchAndReject])
+                {
+                    [expectations removeObject:stub];
+                    [stubs removeObject:stub];
+                }
             }
         }
     }
-    if(removeStub)
-    {
-        @synchronized(stubs)
-        {
-            [stubs removeObject:stub];
-        }
-    }
 
     @try
     {
diff --git a/Source/OCMockTests/OCMockObjectTests.m b/Source/OCMockTests/OCMockObjectTests.m
index eb30a0c4..17505884 100644
--- a/Source/OCMockTests/OCMockObjectTests.m
+++ b/Source/OCMockTests/OCMockObjectTests.m
@@ -1156,4 +1156,23 @@ - (void)testTryingToCreateAnInstanceOfOCMockObjectRaisesAnException
     XCTAssertThrows([[OCMockObject alloc] init]);
 }
 
+#pragma mark thread safety
+
+- (void)testExpectationsAreThreadSafe
+{
+    size_t concurrentTaskCount = 10000;
+
+    for (size_t i = 0; i < concurrentTaskCount; i++)
+    {
+        [[mock expect] lowercaseString];
+    }
+
+    dispatch_apply(concurrentTaskCount, dispatch_get_global_queue(QOS_CLASS_USER_INTERACTIVE, 0), ^(size_t iteration)
+    {
+        [mock lowercaseString];
+    });
+
+    [mock verify];
+}
+
 @end