Port cl/361876089: JniRunnable: Fix a deadlock when Detach() is called from Run() (#330)

Author: dconeybe

firestore/src/android/jni_runnable_android.cc

@@ -1,5 +1,6 @@
 #include "firestore/src/android/jni_runnable_android.h"
 
+#include "app/src/assert.h"
 #include "app/src/util_android.h"
 #include "firestore/src/jni/declaration.h"
 #include "firestore/src/jni/env.h"
@@ -30,9 +31,7 @@ Method<Task> kRunOnNewThread("runOnNewThread",
 Constructor<Object> kConstructor("(J)V");
 
 void NativeRun(JNIEnv* env, jobject java_object, jlong data) {
-  if (data == 0) {
-    return;
-  }
+  FIREBASE_ASSERT_MESSAGE(data != 0, "NativeRun() invoked with data==0");
   reinterpret_cast<JniRunnableBase*>(data)->Run();
 }
 

firestore/src/android/jni_runnable_android.h

@@ -52,12 +52,11 @@ class JniRunnableBase {
    * object's `run()` method will do nothing and complete as if successful.
    *
    * This method will block until all active invocations of `Run()` have
-   * completed, and will cause new invocations of the Java `Runnable` object's
-   * `run()` that occur while this method is blocked to also block until this
-   * method completes.
+   * completed.
    *
-   * Calling `Detach()` multiple times is allowed, but invocations after the
-   * first invocation have no effect.
+   * This method may be safely invoked multiple times. Subsequent invocations
+   * have no side effects but will still block while there are active
+   * invocations of `Run()`.
    */
   void Detach(jni::Env& env);
 
@@ -98,7 +97,8 @@ class JniRunnableBase {
  * A proxy for a Java `Runnable` that calls a C++ function.
  *
  * The template parameter `CallbackT` is typically a lambda or function pointer;
- * it can be anything that can be "invoked" with zero arguments.
+ * it can be anything that can be "invoked" with either zero arguments or one
+ * argument whose type is `JniRunnableBase&`.
  *
  * Example:
  *
@@ -118,9 +118,25 @@ class JniRunnable : public JniRunnableBase {
   JniRunnable(jni::Env& env, CallbackT callback)
       : JniRunnableBase(env), callback_(firebase::Move(callback)) {}
 
-  void Run() override { callback_(); }
+  void Run() override { Run(*this, callback_); }
 
  private:
+  // These two static overloads of `Run()` use SFINAE to invoke the callback
+  // with zero arguments or with one argument, depending on the signature of the
+  // callback. If the callback takes one argument then a reference to the
+  // `JniRunnable` object is specified for that argument.
+  template <typename JniRunnableType, typename ZeroArgCallback>
+  static auto Run(JniRunnableType&, ZeroArgCallback callback)
+      -> decltype(callback()) {
+    return callback();
+  }
+
+  template <typename JniRunnableType, typename OneArgCallback>
+  static auto Run(JniRunnableType& runnable, OneArgCallback callback)
+      -> decltype(callback(runnable)) {
+    return callback(runnable);
+  }
+
   CallbackT callback_;
 };
 

firestore/src/tests/android/jni_runnable_android_test.cc

@@ -1,5 +1,7 @@
 #include "firestore/src/android/jni_runnable_android.h"
 
+#include "app/memory/atomic.h"
+#include "app/src/mutex.h"
 #include "firestore/src/jni/declaration.h"
 #include "firestore/src/jni/object.h"
 #include "firestore/src/jni/ownership.h"
@@ -18,6 +20,7 @@ using jni::Global;
 using jni::Local;
 using jni::Method;
 using jni::Object;
+using jni::StaticField;
 using jni::StaticMethod;
 using jni::Task;
 using jni::Throwable;
@@ -27,14 +30,18 @@ Method<Object> kLooperGetThread("getThread", "()Ljava/lang/Thread;");
 Method<void> kRunnableRun("run", "()V");
 StaticMethod<Object> kCurrentThread("currentThread", "()Ljava/lang/Thread;");
 Method<jlong> kThreadGetId("getId", "()J");
+Method<Object> kThreadGetState("getState", "()Ljava/lang/Thread$State;");
+StaticField<Object> kThreadStateBlocked("BLOCKED", "Ljava/lang/Thread$State;");
 
 class JniRunnableTest : public FirestoreAndroidIntegrationTest {
  public:
   void SetUp() override {
     FirestoreAndroidIntegrationTest::SetUp();
     loader().LoadClass("android/os/Looper", kGetMainLooper, kLooperGetThread);
     loader().LoadClass("java/lang/Runnable", kRunnableRun);
-    loader().LoadClass("java/lang/Thread", kCurrentThread, kThreadGetId);
+    loader().LoadClass("java/lang/Thread", kCurrentThread, kThreadGetId,
+                       kThreadGetState);
+    loader().LoadClass("java/lang/Thread$State", kThreadStateBlocked);
     ASSERT_TRUE(loader().ok());
   }
 };
@@ -56,6 +63,16 @@ jlong GetMainThreadId(Env& env) {
   return env.Call(main_thread, kThreadGetId);
 }
 
+/**
+ * Returns whether or not the given thread is in the "blocked" state.
+ * See java.lang.Thread.State.BLOCKED.
+ */
+bool IsThreadBlocked(Env& env, Object& thread) {
+  Local<Object> actual_state = env.Call(thread, kThreadGetState);
+  Local<Object> expected_state = env.Get(kThreadStateBlocked);
+  return Object::Equals(env, expected_state, actual_state);
+}
+
 TEST_F(JniRunnableTest, JavaRunCallsCppRun) {
   Env env;
   bool invoked = false;
@@ -145,6 +162,27 @@ TEST_F(JniRunnableTest, DetachDetachesEvenIfAnExceptionIsPending) {
   EXPECT_TRUE(env.ok());
 }
 
+// Verify that b/181129657 does not regress; that is, calling `Detach()` from
+// `Run()` should not deadlock.
+TEST_F(JniRunnableTest, DetachCanBeCalledFromRun) {
+  Env env;
+  int run_count = 0;
+  auto runnable = MakeJniRunnable(env, [&run_count](JniRunnableBase& runnable) {
+    ++run_count;
+    Env env;
+    runnable.Detach(env);
+  });
+  Local<Object> java_runnable = runnable.GetJavaRunnable();
+
+  // Call `run()` twice to verify that the call to `Detach()` successfully
+  // detaches and the second `run()` invocation does not call C++ `Run()`.
+  env.Call(java_runnable, kRunnableRun);
+  env.Call(java_runnable, kRunnableRun);
+
+  EXPECT_TRUE(env.ok());
+  EXPECT_EQ(run_count, 1);
+}
+
 TEST_F(JniRunnableTest, DestructionCausesJavaRunToDoNothing) {
   Env env;
   bool invoked = false;
@@ -191,29 +229,21 @@ TEST_F(JniRunnableTest, RunOnMainThreadTaskFailsIfRunThrowsException) {
 }
 
 TEST_F(JniRunnableTest, RunOnMainThreadRunsSynchronouslyFromMainThread) {
-  class ChainedMainThreadJniRunnable : public JniRunnableBase {
-   public:
-    using JniRunnableBase::JniRunnableBase;
-
-    void Run() override {
-      Env env;
-      EXPECT_EQ(GetCurrentThreadId(env), GetMainThreadId(env));
-      if (is_nested_call_) {
-        return;
-      }
-      is_nested_call_ = true;
-      Local<Task> task = RunOnMainThread(env);
-      EXPECT_TRUE(task.IsComplete(env));
-      EXPECT_TRUE(task.IsSuccessful(env));
-      is_nested_call_ = false;
-    }
-
-   private:
-    bool is_nested_call_ = false;
-  };
-
   Env env;
-  ChainedMainThreadJniRunnable runnable(env);
+  bool is_recursive_call = false;
+  auto runnable =
+      MakeJniRunnable(env, [&is_recursive_call](JniRunnableBase& runnable) {
+        Env env;
+        EXPECT_EQ(GetCurrentThreadId(env), GetMainThreadId(env));
+        if (is_recursive_call) {
+          return;
+        }
+        is_recursive_call = true;
+        Local<Task> task = runnable.RunOnMainThread(env);
+        EXPECT_TRUE(task.IsComplete(env));
+        EXPECT_TRUE(task.IsSuccessful(env));
+        is_recursive_call = false;
+      });
 
   Local<Task> task = runnable.RunOnMainThread(env);
 
@@ -252,6 +282,59 @@ TEST_F(JniRunnableTest, RunOnNewThreadTaskFailsIfRunThrowsException) {
   EXPECT_TRUE(env.IsSameObject(exception, thrown_exception));
 }
 
+TEST_F(JniRunnableTest, DetachReturnsAfterLastRunOnAnotherThreadCompletes) {
+  Env env;
+  compat::Atomic<int32_t> runnable1_run_invoke_count;
+  runnable1_run_invoke_count.store(0);
+  Mutex detach_thread_mutex;
+  Global<Object> detach_thread;
+
+  auto runnable1 = MakeJniRunnable(
+      env, [&runnable1_run_invoke_count, &detach_thread, &detach_thread_mutex] {
+        runnable1_run_invoke_count.fetch_add(1);
+        Env env;
+        // Wait for `detach()` to be called and start blocking; then, return to
+        // allow `detach()` to unblock and do its job.
+        while (env.ok()) {
+          MutexLock lock(detach_thread_mutex);
+          if (detach_thread && IsThreadBlocked(env, detach_thread)) {
+            break;
+          }
+        }
+        EXPECT_TRUE(env.ok()) << "IsThreadBlocked() failed with an exception";
+      });
+
+  auto runnable2 =
+      MakeJniRunnable(env, [&runnable1, &detach_thread, &detach_thread_mutex] {
+        Env env;
+        {
+          MutexLock lock(detach_thread_mutex);
+          detach_thread = env.Call(kCurrentThread);
+        }
+        runnable1.Detach(env);
+        EXPECT_TRUE(env.ok()) << "Detach() failed with an exception";
+      });
+
+  // Wait for the `runnable1.Run()` to start to ensure that the lock is held.
+  Local<Task> task1 = runnable1.RunOnNewThread(env);
+  while (true) {
+    if (runnable1_run_invoke_count.load() != 0) {
+      break;
+    }
+  }
+
+  // Start a new thread to call `runnable1.Detach()`.
+  Local<Task> task2 = runnable2.RunOnNewThread(env);
+
+  Await(env, task1);
+  Await(env, task2);
+
+  // Invoke `run()` again and ensure that `Detach()` successfully did its job;
+  // that is, verify that `Run()` is not invoked.
+  env.Call(runnable1.GetJavaRunnable(), kRunnableRun);
+  EXPECT_EQ(runnable1_run_invoke_count.load(), 1);
+}
+
 }  // namespace
 }  // namespace firestore
 }  // namespace firebase

firestore/src_java/com/google/firebase/firestore/internal/cpp/JniRunnable.java

@@ -4,14 +4,11 @@
 import android.os.Looper;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
-import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 /** A {@link Runnable} whose {@link #run} method calls a native function. */
 public final class JniRunnable implements Runnable {
 
-  private final ReentrantReadWriteLock.ReadLock readLock;
-  private final ReentrantReadWriteLock.WriteLock writeLock;
-
+  private final Object lock = new Object();
   private long data;
 
   /**
@@ -26,29 +23,26 @@ public JniRunnable(long data) {
           "data==0 is forbidden because 0 is reserved to indicate that we are detached from the"
               + " C++ function");
     }
-    ReentrantReadWriteLock lock = new ReentrantReadWriteLock(/* fair= */ true);
-    readLock = lock.readLock();
-    writeLock = lock.writeLock();
     this.data = data;
   }
 
   /**
-   * Invokes the C++ method encapsulated by this object.
+   * Invokes the C++ function encapsulated by this object.
    *
    * <p>If {@link #detach} has been invoked then this method does nothing and returns as if
    * successful.
-   *
-   * <p>This method <em>will</em> block if there is a thread blocked in {@link #detach}; otherwise,
-   * it will call the C++ function without blocking. This may even result in concurrent/parallel
-   * calls to the C++ function if {@link #run} is invoked concurrently.
    */
   @Override
   public void run() {
-    readLock.lock();
-    try {
+    // NOTE: Because of the `synchronized` block below, the native function will not be called
+    // concurrently. If concurrent invocations are desired, then this class can be modified with a
+    // more complicated synchronization mechanism.
+    // e.g. https://gist.github.com/dconeybe/2d95fbc75f88de58a49804df5c55157b
+    synchronized (lock) {
+      if (data == 0) {
+        return;
+      }
       nativeRun(data);
-    } finally {
-      readLock.unlock();
     }
   }
 
@@ -58,18 +52,16 @@ public void run() {
    * <p>After this method returns, all future invocations of {@link #run} will do nothing and return
    * as if successful.
    *
-   * <p>This method <em>will</em> block if there are active invocations of {@link #run}. Once all
-   * active invocations of {@link #run} have completed, then this method will proceed and return
-   * nearly instantly. Any invocations of {@link #run} that occur while {@link #detach} is blocked
-   * will also block, allowing the number of active invocations of {@link #run} to eventually reach
-   * zero and allow this method to proceed.
+   * <p>This method blocks until all invocations of the native function called from {@link #run}
+   * complete; therefore, when this method returns it is safe to delete any data that would be
+   * referenced by the native function.
+   *
+   * <p>This method may be safely invoked multiple times. Subsequent invocations have no side
+   * effects but will still block while there are active invocations of the native function.
    */
   public void detach() {
-    writeLock.lock();
-    try {
+    synchronized (lock) {
       data = 0;
-    } finally {
-      writeLock.unlock();
     }
   }