TM iOS: Properly clean up CallbackWrapper during reload

Summary: There is a race condition between tearing down the bridge vs CallbackWrapper invoking the wrapped jsi::Function. This means the wrapper can be stale and unsafe to access after the bridge dies. To avoid unsafe access, let's clean it up properly using weak ref.

Reviewed By: RSNara

Differential Revision: D17380360

fbshipit-source-id: f91ce75d945bf8ed0b141c593bcc674ff465aa8c

Author: fkgozali

ReactCommon/turbomodule/core/TurboModuleBinding.cpp

@@ -44,13 +44,7 @@ void TurboModuleBinding::install(
 }
 
 void TurboModuleBinding::invalidate() const {
-  // TODO (T45804587): Revisit this invalidation logic.
-  // The issue was that Promise resolve/reject functions that are invoked in
-  // some distance future might end up accessing PromiseWrapper that was already
-  // destroyed, if the binding invalidation removed it from the
-  // LongLivedObjectCollection.
-
-  // LongLivedObjectCollection::get().clear();
+  LongLivedObjectCollection::get().clear();
 }
 
 std::shared_ptr<TurboModule> TurboModuleBinding::getModule(

ReactCommon/turbomodule/core/TurboModuleUtils.h

@@ -14,6 +14,7 @@
 #include <jsi/jsi.h>
 
 #include <ReactCommon/JSCallInvoker.h>
+#include <ReactCommon/LongLivedObject.h>
 
 using namespace facebook;
 
@@ -23,7 +24,7 @@ namespace react {
 jsi::Object deepCopyJSIObject(jsi::Runtime &rt, const jsi::Object &obj);
 jsi::Array deepCopyJSIArray(jsi::Runtime &rt, const jsi::Array &arr);
 
-struct Promise {
+struct Promise : public LongLivedObject {
   Promise(jsi::Runtime &rt, jsi::Function resolve, jsi::Function reject);
 
   void resolve(const jsi::Value &result);
@@ -41,7 +42,8 @@ jsi::Value createPromiseAsJSIValue(
     const PromiseSetupFunctionType func);
 
 // Helper for passing jsi::Function arg to other methods.
-class CallbackWrapper {
+// TODO (ramanpreet): Simplify with weak_ptr<>
+class CallbackWrapper : public LongLivedObject {
  private:
   struct Data {
     Data(
@@ -60,6 +62,16 @@ class CallbackWrapper {
   folly::Optional<Data> data_;
 
  public:
+  static std::weak_ptr<CallbackWrapper> createWeak(
+      jsi::Function callback,
+      jsi::Runtime &runtime,
+      std::shared_ptr<react::JSCallInvoker> jsInvoker) {
+    auto wrapper = std::make_shared<CallbackWrapper>(
+        std::move(callback), runtime, jsInvoker);
+    LongLivedObjectCollection::get().add(wrapper);
+    return wrapper;
+  }
+
   CallbackWrapper(
       jsi::Function callback,
       jsi::Runtime &runtime,
@@ -69,6 +81,7 @@ class CallbackWrapper {
   // Delete the enclosed jsi::Function
   void destroy() {
     data_ = folly::none;
+    allowRelease();
   }
 
   bool isDestroyed() {

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

@@ -173,22 +173,28 @@ static RCTResponseSenderBlock convertJSIFunctionToCallback(
     const jsi::Function &value,
     std::shared_ptr<react::JSCallInvoker> jsInvoker)
 {
-  __block auto wrapper = std::make_shared<react::CallbackWrapper>(value.getFunction(runtime), runtime, jsInvoker);
+  auto weakWrapper = react::CallbackWrapper::createWeak(value.getFunction(runtime), runtime, jsInvoker);
+  BOOL __block wrapperWasCalled = NO;
   return ^(NSArray *responses) {
-    if (wrapper == nullptr) {
+    if (wrapperWasCalled) {
       throw std::runtime_error("callback arg cannot be called more than once");
     }
 
-    std::shared_ptr<react::CallbackWrapper> rw = wrapper;
-    wrapper->jsInvoker().invokeAsync([rw, responses]() {
-      std::vector<jsi::Value> args = convertNSArrayToStdVector(rw->runtime(), responses);
-      rw->callback().call(rw->runtime(), (const jsi::Value *)args.data(), args.size());
-    });
+    auto strongWrapper = weakWrapper.lock();
+    if (!strongWrapper) {
+      return;
+    }
 
-    // The callback is single-use, so force release it here.
-    // Doing this also releases the jsi::jsi::Function early, since this block may not get released by ARC for a while,
-    // because the method invocation isn't guarded with @autoreleasepool.
-    wrapper = nullptr;
+    strongWrapper->jsInvoker().invokeAsync([weakWrapper, responses]() {
+      auto strongWrapper2 = weakWrapper.lock();
+      if (!strongWrapper2) {
+        return;
+      }
+
+      std::vector<jsi::Value> args = convertNSArrayToStdVector(strongWrapper2->runtime(), responses);
+      strongWrapper2->callback().call(strongWrapper2->runtime(), (const jsi::Value *)args.data(), args.size());
+      strongWrapper2->destroy();
+    });
   };
 }
 
@@ -222,13 +228,13 @@ static RCTResponseSenderBlock convertJSIFunctionToCallback(
           return jsi::Value::undefined();
         }
 
-        std::shared_ptr<CallbackWrapper> resolveWrapper =
-            std::make_shared<react::CallbackWrapper>(args[0].getObject(rt).getFunction(rt), rt, jsInvoker);
-        std::shared_ptr<CallbackWrapper> rejectWrapper =
-            std::make_shared<react::CallbackWrapper>(args[1].getObject(rt).getFunction(rt), rt, jsInvoker);
+        auto weakResolveWrapper =
+            react::CallbackWrapper::createWeak(args[0].getObject(rt).getFunction(rt), rt, jsInvoker);
+        auto weakRejectWrapper =
+            react::CallbackWrapper::createWeak(args[1].getObject(rt).getFunction(rt), rt, jsInvoker);
 
-        BOOL __block resolveWasCalled = NO;
-        BOOL __block rejectWasCalled = NO;
+        __block BOOL resolveWasCalled = NO;
+        __block BOOL rejectWasCalled = NO;
 
         RCTPromiseResolveBlock resolveBlock = ^(id result) {
           if (rejectWasCalled) {
@@ -239,23 +245,25 @@ static RCTResponseSenderBlock convertJSIFunctionToCallback(
             throw std::runtime_error("Tried to resolve a promise more than once.");
           }
 
-          // In the case that ObjC runtime first invokes this block after
-          // the TurboModuleManager was invalidated, we should do nothing.
-          if (resolveWrapper->isDestroyed()) {
+          auto strongResolveWrapper = weakResolveWrapper.lock();
+          auto strongRejectWrapper = weakRejectWrapper.lock();
+          if (!strongResolveWrapper || !strongRejectWrapper) {
             return;
           }
 
-          resolveWrapper->jsInvoker().invokeAsync([resolveWrapper, rejectWrapper, result]() {
-            if (resolveWrapper->isDestroyed()) {
+          strongResolveWrapper->jsInvoker().invokeAsync([weakResolveWrapper, weakRejectWrapper, result]() {
+            auto strongResolveWrapper2 = weakResolveWrapper.lock();
+            auto strongRejectWrapper2 = weakRejectWrapper.lock();
+            if (!strongResolveWrapper2 || !strongRejectWrapper2) {
               return;
             }
 
-            jsi::Runtime &rt = resolveWrapper->runtime();
+            jsi::Runtime &rt = strongResolveWrapper2->runtime();
             jsi::Value arg = convertObjCObjectToJSIValue(rt, result);
-            resolveWrapper->callback().call(rt, arg);
+            strongResolveWrapper2->callback().call(rt, arg);
 
-            resolveWrapper->destroy();
-            rejectWrapper->destroy();
+            strongResolveWrapper2->destroy();
+            strongRejectWrapper2->destroy();
           });
 
           resolveWasCalled = YES;
@@ -270,24 +278,26 @@ static RCTResponseSenderBlock convertJSIFunctionToCallback(
             throw std::runtime_error("Tried to reject a promise more than once.");
           }
 
-          // In the case that ObjC runtime first invokes this block after
-          // the TurboModuleManager was invalidated, we should do nothing.
-          if (rejectWrapper->isDestroyed()) {
+          auto strongResolveWrapper = weakResolveWrapper.lock();
+          auto strongRejectWrapper = weakRejectWrapper.lock();
+          if (!strongResolveWrapper || !strongRejectWrapper) {
             return;
           }
 
           NSDictionary *jsError = RCTJSErrorFromCodeMessageAndNSError(code, message, error);
-          rejectWrapper->jsInvoker().invokeAsync([rejectWrapper, resolveWrapper, jsError]() {
-            if (rejectWrapper->isDestroyed()) {
+          strongRejectWrapper->jsInvoker().invokeAsync([weakResolveWrapper, weakRejectWrapper, jsError]() {
+            auto strongResolveWrapper2 = weakResolveWrapper.lock();
+            auto strongRejectWrapper2 = weakRejectWrapper.lock();
+            if (!strongResolveWrapper2 || !strongRejectWrapper2) {
               return;
             }
 
-            jsi::Runtime &rt = rejectWrapper->runtime();
+            jsi::Runtime &rt = strongRejectWrapper2->runtime();
             jsi::Value arg = convertNSDictionaryToJSIObject(rt, jsError);
-            rejectWrapper->callback().call(rt, arg);
+            strongRejectWrapper2->callback().call(rt, arg);
 
-            rejectWrapper->destroy();
-            resolveWrapper->destroy();
+            strongResolveWrapper2->destroy();
+            strongRejectWrapper2->destroy();
           });
 
           rejectWasCalled = YES;