Clear all held jsi::Functions when jsi::Runtime is deleted

Summary:
After D19565499, the `LongLivedObjectCollection` will be cleared on the JS thread when the jsi::Runtime is deleted. This diff makes it so that we never hold strong references to `CallbackWrapper`s in our Android TurboModules infra. Therefore, we can leverage the changes in D19565499 to ensure that our `jsi::Function`s are deleted before the `jsi::Runtime`.

## Caveat
If you delete a TurboModule by itself, it's jsi::Functions that haven't been invoked won't be released. This is also the case for iOS. I plan to fix this for both iOS and Android at a later point in time.

Changelog:
[Android][Fixed] - Refactor jsi::Function cleanup in TurboModules

Reviewed By: mdvacca

Differential Revision: D19589151

fbshipit-source-id: efa3cc6c83634014159ac7500dcf6bef9c925762

ReactCommon/turbomodule/core/TurboCxxModule.cpp

@@ -18,23 +18,44 @@ using namespace facebook::xplat::module;
 namespace facebook {
 namespace react {
 
-static CxxModule::Callback makeTurboCxxModuleCallback(
+namespace {
+CxxModule::Callback makeTurboCxxModuleCallback(
     jsi::Runtime &runtime,
-    std::shared_ptr<CallbackWrapper> callbackWrapper) {
-  return [callbackWrapper](std::vector<folly::dynamic> args) {
-    callbackWrapper->jsInvoker().invokeAsync([callbackWrapper, args]() {
+    std::weak_ptr<CallbackWrapper> weakWrapper) {
+  return [weakWrapper,
+          wrapperWasCalled = false](std::vector<folly::dynamic> args) mutable {
+    if (wrapperWasCalled) {
+      throw std::runtime_error("callback arg cannot be called more than once");
+    }
+
+    auto strongWrapper = weakWrapper.lock();
+    if (!strongWrapper) {
+      return;
+    }
+
+    strongWrapper->jsInvoker().invokeAsync([weakWrapper, args]() {
+      auto strongWrapper2 = weakWrapper.lock();
+      if (!strongWrapper2) {
+        return;
+      }
+
       std::vector<jsi::Value> innerArgs;
       for (auto &a : args) {
         innerArgs.push_back(
-            jsi::valueFromDynamic(callbackWrapper->runtime(), a));
+            jsi::valueFromDynamic(strongWrapper2->runtime(), a));
       }
-      callbackWrapper->callback().call(
-          callbackWrapper->runtime(),
+      strongWrapper2->callback().call(
+          strongWrapper2->runtime(),
           (const jsi::Value *)innerArgs.data(),
           innerArgs.size());
+
+      strongWrapper2->destroy();
     });
+
+    wrapperWasCalled = true;
   };
 }
+} // namespace
 
 TurboCxxModule::TurboCxxModule(
     std::unique_ptr<CxxModule> cxxModule,
@@ -132,17 +153,17 @@ jsi::Value TurboCxxModule::invokeMethod(
     }
 
     if (method.callbacks == 1) {
-      auto wrapper = std::make_shared<CallbackWrapper>(
+      auto wrapper = CallbackWrapper::createWeak(
           args[count - 1].getObject(runtime).getFunction(runtime),
           runtime,
           jsInvoker_);
       first = makeTurboCxxModuleCallback(runtime, wrapper);
     } else if (method.callbacks == 2) {
-      auto wrapper1 = std::make_shared<CallbackWrapper>(
+      auto wrapper1 = CallbackWrapper::createWeak(
           args[count - 2].getObject(runtime).getFunction(runtime),
           runtime,
           jsInvoker_);
-      auto wrapper2 = std::make_shared<CallbackWrapper>(
+      auto wrapper2 = CallbackWrapper::createWeak(
           args[count - 1].getObject(runtime).getFunction(runtime),
           runtime,
           jsInvoker_);
@@ -161,9 +182,9 @@ jsi::Value TurboCxxModule::invokeMethod(
         runtime,
         [method, args, count, this](
             jsi::Runtime &rt, std::shared_ptr<Promise> promise) {
-          auto resolveWrapper = std::make_shared<CallbackWrapper>(
+          auto resolveWrapper = CallbackWrapper::createWeak(
               promise->resolve_.getFunction(rt), rt, jsInvoker_);
-          auto rejectWrapper = std::make_shared<CallbackWrapper>(
+          auto rejectWrapper = CallbackWrapper::createWeak(
               promise->reject_.getFunction(rt), rt, jsInvoker_);
           CxxModule::Callback resolve =
               makeTurboCxxModuleCallback(rt, resolveWrapper);

ReactCommon/turbomodule/core/TurboModuleUtils.h

@@ -42,65 +42,46 @@ jsi::Value createPromiseAsJSIValue(
     const PromiseSetupFunctionType func);
 
 // Helper for passing jsi::Function arg to other methods.
-// TODO (ramanpreet): Simplify with weak_ptr<>
 class CallbackWrapper : public LongLivedObject {
  private:
-  struct Data {
-    Data(
-        jsi::Function callback,
-        jsi::Runtime &runtime,
-        std::shared_ptr<react::CallInvoker> jsInvoker)
-        : callback(std::move(callback)),
-          runtime(runtime),
-          jsInvoker(std::move(jsInvoker)) {}
+  CallbackWrapper(
+      jsi::Function &&callback,
+      jsi::Runtime &runtime,
+      std::shared_ptr<CallInvoker> jsInvoker)
+      : callback_(std::move(callback)),
+        runtime_(runtime),
+        jsInvoker_(std::move(jsInvoker)) {}
 
-    jsi::Function callback;
-    jsi::Runtime &runtime;
-    std::shared_ptr<react::CallInvoker> jsInvoker;
-  };
-
-  folly::Optional<Data> data_;
+  jsi::Function callback_;
+  jsi::Runtime &runtime_;
+  std::shared_ptr<CallInvoker> jsInvoker_;
 
  public:
   static std::weak_ptr<CallbackWrapper> createWeak(
-      jsi::Function callback,
+      jsi::Function &&callback,
       jsi::Runtime &runtime,
-      std::shared_ptr<react::CallInvoker> jsInvoker) {
-    auto wrapper = std::make_shared<CallbackWrapper>(
-        std::move(callback), runtime, jsInvoker);
+      std::shared_ptr<CallInvoker> jsInvoker) {
+    auto wrapper = std::shared_ptr<CallbackWrapper>(
+        new CallbackWrapper(std::move(callback), runtime, jsInvoker));
     LongLivedObjectCollection::get().add(wrapper);
     return wrapper;
   }
 
-  CallbackWrapper(
-      jsi::Function callback,
-      jsi::Runtime &runtime,
-      std::shared_ptr<react::CallInvoker> jsInvoker)
-      : data_(Data{std::move(callback), runtime, jsInvoker}) {}
-
   // Delete the enclosed jsi::Function
   void destroy() {
-    data_ = folly::none;
     allowRelease();
   }
 
-  bool isDestroyed() {
-    return !data_.hasValue();
-  }
-
   jsi::Function &callback() {
-    assert(!isDestroyed());
-    return data_->callback;
+    return callback_;
   }
 
   jsi::Runtime &runtime() {
-    assert(!isDestroyed());
-    return data_->runtime;
+    return runtime_;
   }
 
-  react::CallInvoker &jsInvoker() {
-    assert(!isDestroyed());
-    return *(data_->jsInvoker);
+  CallInvoker &jsInvoker() {
+    return *(jsInvoker_);
   }
 };
 

ReactCommon/turbomodule/core/platform/android/JavaTurboModule.cpp

@@ -32,71 +32,60 @@ JavaTurboModule::JavaTurboModule(
       instance_(jni::make_global(instance)),
       nativeInvoker_(nativeInvoker) {}
 
-jni::local_ref<JCxxCallbackImpl::JavaPart>
-JavaTurboModule::createJavaCallbackFromJSIFunction(
-    jsi::Function &function,
+namespace {
+
+jni::local_ref<JCxxCallbackImpl::JavaPart> createJavaCallbackFromJSIFunction(
+    jsi::Function &&function,
     jsi::Runtime &rt,
     std::shared_ptr<CallInvoker> jsInvoker) {
-  auto wrapper = std::make_shared<react::CallbackWrapper>(
-      std::move(function), rt, jsInvoker);
-  callbackWrappers_.insert(wrapper);
+  auto weakWrapper =
+      react::CallbackWrapper::createWeak(std::move(function), rt, jsInvoker);
 
-  std::function<void(folly::dynamic)> fn = [this,
-                                            wrapper](folly::dynamic responses) {
-    if (wrapper->isDestroyed()) {
-      throw std::runtime_error("callback arg cannot be called more than once");
-    }
+  std::function<void(folly::dynamic)> fn =
+      [weakWrapper,
+       wrapperWasCalled = false](folly::dynamic responses) mutable {
+        if (wrapperWasCalled) {
+          throw std::runtime_error(
+              "callback 2 arg cannot be called more than once");
+        }
 
-    wrapper->jsInvoker().invokeAsync([this, wrapper, responses]() {
-      if (wrapper->isDestroyed()) {
-        return;
-      }
+        auto strongWrapper = weakWrapper.lock();
+        if (!strongWrapper) {
+          return;
+        }
 
-      // TODO (T43155926) valueFromDynamic already returns a Value array. Don't
-      // iterate again
-      jsi::Value args = jsi::valueFromDynamic(wrapper->runtime(), responses);
-      auto argsArray =
-          args.getObject(wrapper->runtime()).asArray(wrapper->runtime());
-      std::vector<jsi::Value> result;
-      for (size_t i = 0; i < argsArray.size(wrapper->runtime()); i++) {
-        result.emplace_back(
-            wrapper->runtime(),
-            argsArray.getValueAtIndex(wrapper->runtime(), i));
-      }
-      wrapper->callback().call(
-          wrapper->runtime(), (const jsi::Value *)result.data(), result.size());
+        strongWrapper->jsInvoker().invokeAsync([weakWrapper, responses]() {
+          auto strongWrapper2 = weakWrapper.lock();
+          if (!strongWrapper2) {
+            return;
+          }
 
-      /**
-       * Eagerly destroy the jsi::Function since it's already been invoked.
-       * TODO(T48128233) Do we want callbacks to be invoked only once?
-       *
-       * NOTE: ~JavaTurboModule and this function run on the same thread.
-       * If you reach this point, you know that the destructor wasn't run
-       * because the current wrapper wasn't destroyed. Therefore, it's
-       * safe to access callbackWrappers_.
-       */
-      wrapper->destroy();
-      callbackWrappers_.erase(wrapper);
-    });
-  };
+          // TODO (T43155926) valueFromDynamic already returns a Value array.
+          // Don't iterate again
+          jsi::Value args =
+              jsi::valueFromDynamic(strongWrapper2->runtime(), responses);
+          auto argsArray = args.getObject(strongWrapper2->runtime())
+                               .asArray(strongWrapper2->runtime());
+          std::vector<jsi::Value> result;
+          for (size_t i = 0; i < argsArray.size(strongWrapper2->runtime());
+               i++) {
+            result.emplace_back(
+                strongWrapper2->runtime(),
+                argsArray.getValueAtIndex(strongWrapper2->runtime(), i));
+          }
+          strongWrapper2->callback().call(
+              strongWrapper2->runtime(),
+              (const jsi::Value *)result.data(),
+              result.size());
+
+          strongWrapper2->destroy();
+        });
+
+        wrapperWasCalled = true;
+      };
   return JCxxCallbackImpl::newObjectCxxArgs(fn);
 }
 
-JavaTurboModule::~JavaTurboModule() {
-  /**
-   * Delete all jsi::Functions that haven't yet been invoked by Java.
-   * So long as nothing else aside from the JS heap is holding on to this
-   * JavaTurboModule, this destructor is guaranteed to execute before the
-   * jsi::Runtime is deleted.
-   */
-  for (auto it = callbackWrappers_.begin(); it != callbackWrappers_.end();
-       it++) {
-    (*it)->destroy();
-  }
-}
-
-namespace {
-
 template <typename T>
 std::string to_string(T v) {
   std::ostringstream stream;
@@ -362,7 +351,8 @@ JNIArgs JavaTurboModule::convertJSIArgsToJNIArgs(
 
       jsi::Function fn = arg->getObject(rt).getFunction(rt);
       jarg->l = makeGlobalIfNecessary(
-          createJavaCallbackFromJSIFunction(fn, rt, jsInvoker).release());
+          createJavaCallbackFromJSIFunction(std::move(fn), rt, jsInvoker)
+              .release());
       continue;
     }
 
@@ -618,10 +608,10 @@ jsi::Value JavaTurboModule::invokeJavaMethod(
                     runtime);
 
             auto resolve = createJavaCallbackFromJSIFunction(
-                               resolveJSIFn, runtime, jsInvoker_)
+                               std::move(resolveJSIFn), runtime, jsInvoker_)
                                .release();
             auto reject = createJavaCallbackFromJSIFunction(
-                              rejectJSIFn, runtime, jsInvoker_)
+                              std::move(rejectJSIFn), runtime, jsInvoker_)
                               .release();
 
             jclass jPromiseImpl =

ReactCommon/turbomodule/core/platform/android/JavaTurboModule.h

@@ -45,26 +45,10 @@ class JSI_EXPORT JavaTurboModule : public TurboModule {
       const jsi::Value *args,
       size_t argCount);
 
-  /**
-   * This dtor must be called from the JS Thread, since it accesses
-   * callbackWrappers_, which createJavaCallbackFromJSIFunction also accesses
-   * from the JS Thread.
-   */
-  virtual ~JavaTurboModule();
-
  private:
   jni::global_ref<JTurboModule> instance_;
-  std::unordered_set<std::shared_ptr<CallbackWrapper>> callbackWrappers_;
   std::shared_ptr<CallInvoker> nativeInvoker_;
 
-  /**
-   * This method must be called from the JS Thread, since it accesses
-   * callbackWrappers_.
-   */
-  jni::local_ref<JCxxCallbackImpl::JavaPart> createJavaCallbackFromJSIFunction(
-      jsi::Function &function,
-      jsi::Runtime &rt,
-      std::shared_ptr<CallInvoker> jsInvoker);
   JNIArgs convertJSIArgsToJNIArgs(
       JNIEnv *env,
       jsi::Runtime &rt,