WebWorkers: Allow native modules to be notified when executors are unregistered

Summary: This will allow them to clean up resources when a web worker goes away.

Reviewed By: mhorowitz, lexs

Differential Revision: D2994721

fb-gh-sync-id: c7ca1afc7290e85038cf692a139f6478dba0ef61
shipit-source-id: c7ca1afc7290e85038cf692a139f6478dba0ef61

ReactAndroid/src/main/java/com/facebook/react/bridge/CatalystInstanceImpl.java

@@ -222,8 +222,10 @@ public class CatalystInstanceImpl implements CatalystInstance {
       Systrace.unregisterListener(mTraceListener);
 
       synchronouslyDisposeBridgeOnJSThread();
-      mReactQueueConfiguration.destroy();
     }
+
+    mReactQueueConfiguration.destroy();
+
     boolean wasIdle = (mPendingJSCalls.getAndSet(0) == 0);
     if (!wasIdle && !mBridgeIdleListeners.isEmpty()) {
       for (NotThreadSafeBridgeIdleDebugListener listener : mBridgeIdleListeners) {
@@ -293,12 +295,12 @@ public class CatalystInstanceImpl implements CatalystInstance {
   @Override
   public void handleMemoryPressure(final MemoryPressure level) {
     mReactQueueConfiguration.getJSQueueThread().runOnQueue(
-        new Runnable() {
-          @Override
-          public void run() {
-            Assertions.assertNotNull(mBridge).handleMemoryPressure(level);
-          }
-        });
+      new Runnable() {
+        @Override
+        public void run() {
+          Assertions.assertNotNull(mBridge).handleMemoryPressure(level);
+        }
+      });
   }
 
   /**
@@ -399,12 +401,14 @@ public class CatalystInstanceImpl implements CatalystInstance {
     public void call(ExecutorToken executorToken, int moduleId, int methodId, ReadableNativeArray parameters) {
       mReactQueueConfiguration.getNativeModulesQueueThread().assertIsOnThread();
 
-      // Suppress any callbacks if destroyed - will only lead to sadness.
-      if (mDestroyed) {
-        return;
-      }
+      synchronized (mTeardownLock) {
+        // Suppress any callbacks if destroyed - will only lead to sadness.
+        if (mDestroyed) {
+          return;
+        }
 
-      mJavaRegistry.call(CatalystInstanceImpl.this, executorToken, moduleId, methodId, parameters);
+        mJavaRegistry.call(CatalystInstanceImpl.this, executorToken, moduleId, methodId, parameters);
+      }
     }
 
     @Override
@@ -415,17 +419,38 @@ public class CatalystInstanceImpl implements CatalystInstance {
       // native modules could be in a bad state so we don't want to call anything on them. We
       // still want to trigger the debug listener since it allows instrumentation tests to end and
       // check their assertions without waiting for a timeout.
-      if (!mDestroyed) {
-        Systrace.beginSection(Systrace.TRACE_TAG_REACT_JAVA_BRIDGE, "onBatchComplete");
-        try {
-          mJavaRegistry.onBatchComplete();
-        } finally {
-          Systrace.endSection(Systrace.TRACE_TAG_REACT_JAVA_BRIDGE);
+      synchronized (mTeardownLock) {
+        if (!mDestroyed) {
+          Systrace.beginSection(Systrace.TRACE_TAG_REACT_JAVA_BRIDGE, "onBatchComplete");
+          try {
+            mJavaRegistry.onBatchComplete();
+          } finally {
+            Systrace.endSection(Systrace.TRACE_TAG_REACT_JAVA_BRIDGE);
+          }
         }
       }
 
       decrementPendingJSCalls();
     }
+
+    @Override
+    public void onExecutorUnregistered(ExecutorToken executorToken) {
+      mReactQueueConfiguration.getNativeModulesQueueThread().assertIsOnThread();
+
+      // Since onCatalystInstanceDestroy happens on the UI thread, we don't want to also execute
+      // this callback on the native modules thread at the same time. Longer term, onCatalystInstanceDestroy
+      // should probably be executed on the native modules thread as well instead.
+      synchronized (mTeardownLock) {
+        if (!mDestroyed) {
+          Systrace.beginSection(Systrace.TRACE_TAG_REACT_JAVA_BRIDGE, "onExecutorUnregistered");
+          try {
+            mJavaRegistry.onExecutorUnregistered(executorToken);
+          } finally {
+            Systrace.endSection(Systrace.TRACE_TAG_REACT_JAVA_BRIDGE);
+          }
+        }
+      }
+    }
   }
 
   private class NativeExceptionHandler implements QueueThreadExceptionHandler {

ReactAndroid/src/main/java/com/facebook/react/bridge/NativeModuleRegistry.java

@@ -30,6 +30,7 @@ public class NativeModuleRegistry {
   private final List<ModuleDefinition> mModuleTable;
   private final Map<Class<? extends NativeModule>, NativeModule> mModuleInstances;
   private final ArrayList<OnBatchCompleteListener> mBatchCompleteListenerModules;
+  private final ArrayList<OnExecutorUnregisteredListener> mOnExecutorUnregisteredListenerModules;
 
   private NativeModuleRegistry(
       List<ModuleDefinition> moduleTable,
@@ -37,12 +38,16 @@ public class NativeModuleRegistry {
     mModuleTable = moduleTable;
     mModuleInstances = moduleInstances;
 
-    mBatchCompleteListenerModules = new ArrayList<OnBatchCompleteListener>(mModuleTable.size());
+    mBatchCompleteListenerModules = new ArrayList<>(mModuleTable.size());
+    mOnExecutorUnregisteredListenerModules = new ArrayList<>(mModuleTable.size());
     for (int i = 0; i < mModuleTable.size(); i++) {
       ModuleDefinition definition = mModuleTable.get(i);
       if (definition.target instanceof OnBatchCompleteListener) {
         mBatchCompleteListenerModules.add((OnBatchCompleteListener) definition.target);
       }
+      if (definition.target instanceof OnExecutorUnregisteredListener) {
+        mOnExecutorUnregisteredListenerModules.add((OnExecutorUnregisteredListener) definition.target);
+      }
     }
   }
 
@@ -135,6 +140,12 @@ public class NativeModuleRegistry {
     }
   }
 
+  public void onExecutorUnregistered(ExecutorToken executorToken) {
+    for (int i = 0; i < mOnExecutorUnregisteredListenerModules.size(); i++) {
+      mOnExecutorUnregisteredListenerModules.get(i).onExecutorDestroyed(executorToken);
+    }
+  }
+
   public <T extends NativeModule> T getModule(Class<T> moduleInterface) {
     return (T) Assertions.assertNotNull(mModuleInstances.get(moduleInterface));
   }

ReactAndroid/src/main/java/com/facebook/react/bridge/OnExecutorUnregisteredListener.java

@@ -0,0 +1,22 @@
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree. An additional grant
+ * of patent rights can be found in the PATENTS file in the same directory.
+ */
+
+package com.facebook.react.bridge;
+
+/**
+ * Interface for a module that will be notified when JS executors have been unregistered from the bridge.
+ * Note that this will NOT notify listeners about the main executor being destroyed: use
+ * {@link NativeModule#onCatalystInstanceDestroy()} for that. Once a module has received a
+ * {@link NativeModule#onCatalystInstanceDestroy()} call, it will not receive any onExecutorUnregistered
+ * calls.
+ */
+public interface OnExecutorUnregisteredListener {
+
+  void onExecutorDestroyed(ExecutorToken executorToken);
+}

ReactAndroid/src/main/java/com/facebook/react/bridge/ReactCallback.java

@@ -19,4 +19,7 @@ public interface ReactCallback {
 
   @DoNotStrip
   void onBatchComplete();
+
+  @DoNotStrip
+  void onExecutorUnregistered(ExecutorToken executorToken);
 }

ReactAndroid/src/main/jni/react/Bridge.cpp

@@ -17,7 +17,7 @@ namespace react {
 Bridge::Bridge(
     JSExecutorFactory* jsExecutorFactory,
     std::unique_ptr<ExecutorTokenFactory> executorTokenFactory,
-    Callback callback) :
+    std::unique_ptr<BridgeCallback> callback) :
   m_callback(std::move(callback)),
   m_destroyed(std::make_shared<bool>(false)),
   m_executorTokenFactory(std::move(executorTokenFactory)) {
@@ -174,7 +174,7 @@ void Bridge::callNativeModules(JSExecutor& executor, const std::string& callJSON
   if (*m_destroyed) {
     return;
   }
-  m_callback(getTokenForExecutor(executor), parseMethodCalls(callJSON), isEndOfBatch);
+  m_callback->onCallNativeModules(getTokenForExecutor(executor), parseMethodCalls(callJSON), isEndOfBatch);
 }
 
 ExecutorToken Bridge::getMainExecutorToken() const {
@@ -214,7 +214,7 @@ std::unique_ptr<JSExecutor> Bridge::unregisterExecutor(ExecutorToken executorTok
     m_executorTokenMap.erase(executor.get());
   }
 
-  // TODO: Notify native modules that ExecutorToken destroyed
+  m_callback->onExecutorUnregistered(executorToken);
 
   return executor;
 }

ReactAndroid/src/main/jni/react/Bridge.h

@@ -24,6 +24,18 @@ struct dynamic;
 namespace facebook {
 namespace react {
 
+class BridgeCallback {
+public:
+  virtual ~BridgeCallback() {};
+
+  virtual void onCallNativeModules(
+      ExecutorToken executorToken,
+      std::vector<MethodCall>&& calls,
+      bool isEndOfBatch) = 0;
+
+  virtual void onExecutorUnregistered(ExecutorToken executorToken) = 0;
+};
+
 class Bridge;
 class ExecutorRegistration {
 public:
@@ -39,15 +51,13 @@ public:
 
 class Bridge {
 public:
-  typedef std::function<void(ExecutorToken executorToken, std::vector<MethodCall>, bool isEndOfBatch)> Callback;
-
   /**
    * This must be called on the main JS thread.
    */
   Bridge(
       JSExecutorFactory* jsExecutorFactory,
       std::unique_ptr<ExecutorTokenFactory> executorTokenFactory,
-      Callback callback);
+      std::unique_ptr<BridgeCallback> callback);
   virtual ~Bridge();
 
   /**
@@ -128,7 +138,7 @@ public:
    */
   void destroy();
 private:
-  Callback m_callback;
+  std::unique_ptr<BridgeCallback> m_callback;
   // This is used to avoid a race condition where a proxyCallback gets queued after ~Bridge(),
   // on the same thread. In that case, the callback will try to run the task on m_callback which
   // will have been destroyed within ~Bridge(), thus causing a SIGSEGV.

ReactAndroid/src/main/jni/react/jni/OnLoad.cpp

@@ -567,6 +567,7 @@ namespace bridge {
 
 static jmethodID gCallbackMethod;
 static jmethodID gOnBatchCompleteMethod;
+static jmethodID gOnExecutorUnregisteredMethod;
 static jmethodID gLogMarkerMethod;
 
 struct CountableBridge : Bridge, Countable {
@@ -582,7 +583,7 @@ static void logMarker(const std::string& marker) {
   env->DeleteLocalRef(jmarker);
 }
 
-static void makeJavaCall(JNIEnv* env, ExecutorToken executorToken, jobject callback, MethodCall&& call) {
+static void makeJavaCall(JNIEnv* env, ExecutorToken executorToken, jobject callback, const MethodCall& call) {
   if (call.arguments.isNull()) {
     return;
   }
@@ -607,33 +608,51 @@ static void signalBatchComplete(JNIEnv* env, jobject callback) {
   env->CallVoidMethod(callback, gOnBatchCompleteMethod);
 }
 
-static void dispatchCallbacksToJava(ExecutorToken executorToken,
-                                    const RefPtr<WeakReference>& weakCallback,
-                                    const RefPtr<WeakReference>& weakCallbackQueueThread,
-                                    std::vector<MethodCall>&& calls,
-                                    bool isEndOfBatch) {
-  auto env = Environment::current();
-  if (env->ExceptionCheck()) {
-    FBLOGW("Dropped calls because of pending exception");
-    return;
-  }
+class PlatformBridgeCallback : public BridgeCallback {
+public:
+  PlatformBridgeCallback(
+      RefPtr<WeakReference> weakCallback_,
+      RefPtr<WeakReference> weakCallbackQueueThread_) :
+    weakCallback_(std::move(weakCallback_)),
+    weakCallbackQueueThread_(std::move(weakCallbackQueueThread_)) {}
 
-  ResolvedWeakReference callbackQueueThread(weakCallbackQueueThread);
-  if (!callbackQueueThread) {
-    FBLOGW("Dropped calls because of callback queue thread went away");
-    return;
-  }
-
-  auto runnableFunction = std::bind([executorToken, weakCallback, isEndOfBatch] (std::vector<MethodCall>& calls) {
+  void executeCallbackOnCallbackQueueThread(std::function<void(ResolvedWeakReference&)>&& runnable) {
     auto env = Environment::current();
     if (env->ExceptionCheck()) {
-      FBLOGW("Dropped calls because of pending exception");
+      FBLOGW("Dropped callback because of pending exception");
       return;
     }
-    ResolvedWeakReference callback(weakCallback);
-    if (callback) {
-      for (auto&& call : calls) {
-        makeJavaCall(env, executorToken, callback, std::move(call));
+
+    ResolvedWeakReference callbackQueueThread(weakCallbackQueueThread_);
+    if (!callbackQueueThread) {
+      FBLOGW("Dropped callback because callback queue thread went away");
+      return;
+    }
+
+    auto runnableWrapper = std::bind([this] (std::function<void(ResolvedWeakReference&)>& runnable) {
+      auto env = Environment::current();
+      if (env->ExceptionCheck()) {
+        FBLOGW("Dropped calls because of pending exception");
+        return;
+      }
+      ResolvedWeakReference callback(weakCallback_);
+      if (callback) {
+        runnable(callback);
+      }
+    }, std::move(runnable));
+
+    auto jNativeRunnable = runnable::createNativeRunnable(env, std::move(runnableWrapper));
+    queue::enqueueNativeRunnableOnQueue(env, callbackQueueThread, jNativeRunnable.get());
+  }
+
+  virtual void onCallNativeModules(
+      ExecutorToken executorToken,
+      std::vector<MethodCall>&& calls,
+      bool isEndOfBatch) override {
+    executeCallbackOnCallbackQueueThread([executorToken, calls, isEndOfBatch] (ResolvedWeakReference& callback) {
+      JNIEnv* env = Environment::current();
+      for (auto& call : calls) {
+        makeJavaCall(env, executorToken, callback, call);
         if (env->ExceptionCheck()) {
           return;
         }
@@ -641,23 +660,31 @@ static void dispatchCallbacksToJava(ExecutorToken executorToken,
       if (isEndOfBatch) {
         signalBatchComplete(env, callback);
       }
-    }
-  }, std::move(calls));
+    });
+  }
 
-  auto jNativeRunnable = runnable::createNativeRunnable(env, std::move(runnableFunction));
-  queue::enqueueNativeRunnableOnQueue(env, callbackQueueThread, jNativeRunnable.get());
-}
+  virtual void onExecutorUnregistered(ExecutorToken executorToken) override {
+    executeCallbackOnCallbackQueueThread([executorToken] (ResolvedWeakReference& callback) {
+      JNIEnv *env = Environment::current();
+      env->CallVoidMethod(
+          callback,
+          gOnExecutorUnregisteredMethod,
+          static_cast<JExecutorTokenHolder*>(executorToken.getPlatformExecutorToken().get())->getJobj());
+    });
+  }
+private:
+  RefPtr<WeakReference> weakCallback_;
+  RefPtr<WeakReference> weakCallbackQueueThread_;
+};
 
 static void create(JNIEnv* env, jobject obj, jobject executor, jobject callback,
                    jobject callbackQueueThread) {
   auto weakCallback = createNew<WeakReference>(callback);
   auto weakCallbackQueueThread = createNew<WeakReference>(callbackQueueThread);
-  auto bridgeCallback = [weakCallback, weakCallbackQueueThread] (ExecutorToken executorToken, std::vector<MethodCall> calls, bool isEndOfBatch) {
-    dispatchCallbacksToJava(executorToken, weakCallback, weakCallbackQueueThread, std::move(calls), isEndOfBatch);
-  };
+  auto bridgeCallback = folly::make_unique<PlatformBridgeCallback>(weakCallback, weakCallbackQueueThread);
   auto nativeExecutorFactory = extractRefPtr<CountableJSExecutorFactory>(env, executor);
   auto executorTokenFactory = folly::make_unique<JExecutorTokenFactory>();
-  auto bridge = createNew<CountableBridge>(nativeExecutorFactory.get(), std::move(executorTokenFactory), bridgeCallback);
+  auto bridge = createNew<CountableBridge>(nativeExecutorFactory.get(), std::move(executorTokenFactory), std::move(bridgeCallback));
   setCountableForJava(env, obj, std::move(bridge));
 }
 
@@ -968,6 +995,7 @@ extern "C" JNIEXPORT jint JNI_OnLoad(JavaVM* vm, void* reserved) {
     jclass callbackClass = env->FindClass("com/facebook/react/bridge/ReactCallback");
     bridge::gCallbackMethod = env->GetMethodID(callbackClass, "call", "(Lcom/facebook/react/bridge/ExecutorToken;IILcom/facebook/react/bridge/ReadableNativeArray;)V");
     bridge::gOnBatchCompleteMethod = env->GetMethodID(callbackClass, "onBatchComplete", "()V");
+    bridge::gOnExecutorUnregisteredMethod = env->GetMethodID(callbackClass, "onExecutorUnregistered", "(Lcom/facebook/react/bridge/ExecutorToken;)V");
 
     jclass markerClass = env->FindClass("com/facebook/react/bridge/ReactMarker");
     bridge::gLogMarkerMethod = env->GetStaticMethodID(markerClass, "logMarker", "(Ljava/lang/String;)V");