Fix crash in AsyncCallback::callWithFunction (#41001)

Summary:
Pull Request resolved: https://github.com/facebook/react-native/pull/41001

I believe this is due to a race condition between VM teardown and callback invocation. Because we were previously retaining the CallbackWrapper across the invokeAsync call, we may potentially have been holding onto the JSI::Function after it was already destroyed.

Changelog: [Internal]

Reviewed By: sammy-SC

Differential Revision: D50286876

fbshipit-source-id: 1277a9f37166da59ebb2169fe8d5a6fabce82f1b

packages/react-native/ReactCommon/react/bridging/CallbackWrapper.h

@@ -8,10 +8,13 @@
 #pragma once
 
 #include <jsi/jsi.h>
-#include "LongLivedObject.h"
 
 #include <memory>
 
+#include <ReactCommon/CallInvoker.h>
+
+#include "LongLivedObject.h"
+
 namespace facebook::react {
 
 // Helper for passing jsi::Function arg to other methods.

packages/react-native/ReactCommon/react/bridging/Function.h

@@ -61,15 +61,14 @@ class AsyncCallback {
 
   void callWithArgs(std::optional<SchedulerPriority> priority, Args... args)
       const noexcept {
-    auto wrapper = callback_->wrapper_.lock();
-    if (wrapper) {
-      auto& jsInvoker = wrapper->jsInvoker();
+    if (auto wrapper = callback_->wrapper_.lock()) {
       auto fn = [callback = callback_,
                  argsPtr = std::make_shared<std::tuple<Args...>>(
                      std::make_tuple(std::forward<Args>(args)...))] {
         callback->apply(std::move(*argsPtr));
       };
 
+      auto& jsInvoker = wrapper->jsInvoker();
       if (priority) {
         jsInvoker.invokeAsync(*priority, std::move(fn));
       } else {
@@ -82,14 +81,17 @@ class AsyncCallback {
       std::optional<SchedulerPriority> priority,
       std::function<void(jsi::Runtime&, jsi::Function&)>&& callImpl)
       const noexcept {
-    auto wrapper = callback_->wrapper_.lock();
-    if (wrapper) {
-      auto& jsInvoker = wrapper->jsInvoker();
-      auto fn = [wrapper = std::move(wrapper),
-                 callImpl = std::move(callImpl)]() {
-        callImpl(wrapper->runtime(), wrapper->callback());
+    if (auto wrapper = callback_->wrapper_.lock()) {
+      // Capture callback_ and not wrapper_. If callback_ is deallocated or the
+      // JSVM is shutdown before the async task is scheduled, the underlying
+      // function will have been deallocated.
+      auto fn = [callback = callback_, callImpl = std::move(callImpl)]() {
+        if (auto wrapper2 = callback->wrapper_.lock()) {
+          callImpl(wrapper2->runtime(), wrapper2->callback());
+        }
       };
 
+      auto& jsInvoker = wrapper->jsInvoker();
       if (priority) {
         jsInvoker.invokeAsync(*priority, std::move(fn));
       } else {

packages/react-native/ReactCommon/react/bridging/LongLivedObject.h

@@ -19,9 +19,12 @@ namespace facebook::react {
  * collection when needed.
  *
  * The subclass of this class must be created using std::make_shared<T>().
- * After creation, add it to the `LongLivedObjectCollection`.
- * When done with the object, call `allowRelease()` to allow the OS to release
- * it.
+ * After creation, add it to the `LongLivedObjectCollection`. When done with the
+ * object, call `allowRelease()` to reclaim its memory.
+ *
+ * When using LongLivedObject to keep JS values alive, ensure you only hold weak
+ * references to the object outside the JS thread to avoid accessing deallocated
+ * values when the JS VM is shutdown.
  */
 class LongLivedObject {
  public:

packages/react-native/ReactCommon/react/bridging/tests/BridgingTest.cpp

@@ -290,6 +290,24 @@ TEST_F(BridgingTest, asyncCallbackTest) {
   EXPECT_EQ("hello again"s, output);
 }
 
+TEST_F(BridgingTest, asyncCallbackInvalidation) {
+  std::string output;
+  std::function<void(std::string)> func = [&](auto str) { output = str; };
+
+  auto jsCallback = bridging::fromJs<AsyncCallback<>>(
+      rt, bridging::toJs(rt, func, invoker), invoker);
+  jsCallback.call(
+      [](jsi::Runtime& rt, jsi::Function& f) { f.call(rt, "hello"); });
+
+  // LongLivedObjectCollection goes away before callback is executed
+  LongLivedObjectCollection::get().clear();
+
+  flushQueue();
+
+  // Assert native callback is never invoked
+  ASSERT_EQ(""s, output);
+}
+
 TEST_F(BridgingTest, asyncCallbackImplicitBridgingTest) {
   std::string output;
   auto func = std::function<void(std::string)>([&](auto str) { output = str; });