TM iOS: Prevent module cache invalidation race

Summary: There are still race condition during bridge invalidation. Some modules may be accessing other modules during invalidation, but it's racing with the TM manager clearing the cache.

Reviewed By: JoshuaGross

Differential Revision: D15947488

fbshipit-source-id: 3bd51382264f538a03ca565b0f099da40c3daadf

React/Base/RCTBridge.h

@@ -85,6 +85,12 @@ RCT_EXTERN NSString *const RCTBridgeWillDownloadScriptNotification;
  */
 RCT_EXTERN NSString *const RCTBridgeDidDownloadScriptNotification;
 
+/**
+ * This notification fires right after the bridge is about to invalidate NativeModule
+ * instances during teardown. Handle this notification to perform additional invalidation.
+ */
+RCT_EXTERN NSString *const RCTBridgeWillInvalidateModulesNotification;
+
 /**
  * This notification fires right after the bridge finishes invalidating NativeModule
  * instances during teardown. Handle this notification to perform additional invalidation.

React/Base/RCTBridge.m

@@ -33,6 +33,7 @@ NSString *const RCTDidSetupModuleNotificationSetupTimeKey = @"setupTime";
 NSString *const RCTBridgeWillReloadNotification = @"RCTBridgeWillReloadNotification";
 NSString *const RCTBridgeWillDownloadScriptNotification = @"RCTBridgeWillDownloadScriptNotification";
 NSString *const RCTBridgeDidDownloadScriptNotification = @"RCTBridgeDidDownloadScriptNotification";
+NSString *const RCTBridgeWillInvalidateModulesNotification = @"RCTBridgeWillInvalidateModulesNotification";
 NSString *const RCTBridgeDidInvalidateModulesNotification = @"RCTBridgeDidInvalidateModulesNotification";
 NSString *const RCTBridgeDidDownloadScriptNotificationSourceKey = @"source";
 NSString *const RCTBridgeDidDownloadScriptNotificationBridgeDescriptionKey = @"bridgeDescription";

React/CxxBridge/RCTCxxBridge.mm

@@ -1077,6 +1077,11 @@ RCT_NOT_IMPLEMENTED(- (instancetype)initWithBundleURL:(__unused NSURL *)bundleUR
     }
 
     // Invalidate modules
+
+    [[NSNotificationCenter defaultCenter] postNotificationName:RCTBridgeWillInvalidateModulesNotification
+                                                        object:self->_parentBridge
+                                                      userInfo:@{@"bridge": self}];
+
     // We're on the JS thread (which we'll be suspending soon), so no new calls will be made to native modules after
     // this completes. We must ensure all previous calls were dispatched before deallocating the instance (and module
     // wrappers) or we may have invalid pointers still in flight.
@@ -1098,14 +1103,14 @@ RCT_NOT_IMPLEMENTED(- (instancetype)initWithBundleURL:(__unused NSURL *)bundleUR
       [moduleData invalidate];
     }
 
-    [[NSNotificationCenter defaultCenter] postNotificationName:RCTBridgeDidInvalidateModulesNotification
-                                                        object:self->_parentBridge
-                                                      userInfo:@{@"bridge": self}];
-
     if (dispatch_group_wait(moduleInvalidation, dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC))) {
       RCTLogError(@"Timed out waiting for modules to be invalidated");
     }
 
+    [[NSNotificationCenter defaultCenter] postNotificationName:RCTBridgeDidInvalidateModulesNotification
+                                                        object:self->_parentBridge
+                                                      userInfo:@{@"bridge": self}];
+
     self->_reactInstance.reset();
     self->_jsMessageThread.reset();
 

ReactCommon/turbomodule/core/platform/ios/RCTTurboModuleManager.mm

@@ -7,6 +7,7 @@
 
 #import "RCTTurboModuleManager.h"
 
+#import <atomic>
 #import <cassert>
 #import <mutex>
 
@@ -58,6 +59,7 @@ static Class getFallbackClassFromName(const char *name)
    * JS thread.
    */
   std::mutex _rctTurboModuleCacheLock;
+  std::atomic<bool> _invalidating;
 }
 
 - (instancetype)initWithBridge:(RCTBridge *)bridge delegate:(id<RCTTurboModuleManagerDelegate>)delegate
@@ -66,10 +68,15 @@ static Class getFallbackClassFromName(const char *name)
     _jsInvoker = std::make_shared<react::BridgeJSCallInvoker>(bridge.reactInstance);
     _delegate = delegate;
     _bridge = bridge;
+    _invalidating = false;
 
     // Necessary to allow NativeModules to lookup TurboModules
     [bridge setRCTTurboModuleLookupDelegate:self];
 
+    [[NSNotificationCenter defaultCenter] addObserver:self
+                                             selector:@selector(bridgeWillInvalidateModules:)
+                                                 name:RCTBridgeWillInvalidateModulesNotification
+                                               object:_bridge.parentBridge];
     [[NSNotificationCenter defaultCenter] addObserver:self
                                              selector:@selector(bridgeDidInvalidateModules:)
                                                  name:RCTBridgeDidInvalidateModulesNotification
@@ -222,6 +229,11 @@ static Class getFallbackClassFromName(const char *name)
       return rctTurboModuleCacheLookup->second;
     }
 
+    if (_invalidating) {
+      // Don't allow creating new instances while invalidating.
+      return nil;
+    }
+
     /**
      * Step 2a: Resolve platform-specific class.
      */
@@ -369,17 +381,32 @@ static Class getFallbackClassFromName(const char *name)
 
 #pragma mark Invalidation logic
 
-- (void)bridgeDidInvalidateModules:(NSNotification *)notification
+- (void)bridgeWillInvalidateModules:(NSNotification *)notification
 {
-  // Rely on this notification to invalidate all known TurboModule ObjC instances, synchronously.
   RCTBridge *bridge = notification.userInfo[@"bridge"];
   if (bridge != _bridge) {
     return;
   }
 
+  _invalidating = true;
+}
+
+- (void)bridgeDidInvalidateModules:(NSNotification *)notification
+{
+  RCTBridge *bridge = notification.userInfo[@"bridge"];
+  if (bridge != _bridge) {
+    return;
+  }
+
+  std::unordered_map<std::string, id<RCTTurboModule>> rctCacheCopy;
+  {
+    std::unique_lock<std::mutex> lock(_rctTurboModuleCacheLock);
+    rctCacheCopy.insert(_rctTurboModuleCache.begin(), _rctTurboModuleCache.end());
+  }
+
   // Backward-compatibility: RCTInvalidating handling.
   dispatch_group_t moduleInvalidationGroup = dispatch_group_create();
-  for (const auto &p : _rctTurboModuleCache) {
+  for (const auto &p : rctCacheCopy) {
     id<RCTTurboModule> module = p.second;
     if ([module respondsToSelector:@selector(invalidate)]) {
       if ([module respondsToSelector:@selector(methodQueue)]) {
@@ -400,10 +427,14 @@ static Class getFallbackClassFromName(const char *name)
   }
 
   if (dispatch_group_wait(moduleInvalidationGroup, dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC))) {
-    RCTLogError(@"Timed out waiting for modules to be invalidated");
+    RCTLogError(@"TurboModuleManager: Timed out waiting for modules to be invalidated");
+  }
+
+  {
+    std::unique_lock<std::mutex> lock(_rctTurboModuleCacheLock);
+    _rctTurboModuleCache.clear();
   }
 
-  _rctTurboModuleCache.clear();
   _turboModuleCache.clear();
 
   _binding->invalidate();