Address feedback

Author: cherylEnkidu

firestore/integration_test_internal/CMakeLists.txt

@@ -121,8 +121,6 @@ set(FIREBASE_INTEGRATION_TEST_PORTABLE_TEST_SRCS
 
 # These sources contain the actual tests that run on Android only.
 set(FIREBASE_INTEGRATION_TEST_ANDROID_TEST_SRCS
-    src/android/arena_ref_android_test.cc
-    src/android/env_android_test.cc
     src/android/field_path_portable_test.cc
     src/android/firestore_integration_test_android_test.cc
     src/android/geo_point_android_test.cc
@@ -133,6 +131,7 @@ set(FIREBASE_INTEGRATION_TEST_ANDROID_TEST_SRCS
     src/android/snapshot_metadata_android_test.cc
     src/android/timestamp_android_test.cc
     src/android/transaction_options_android_test.cc
+    src/jni/arena_ref_test.cc
     src/jni/declaration_test.cc
     src/jni/env_test.cc
     src/jni/object_test.cc

firestore/integration_test_internal/src/android/env_android_test.cc

@@ -14,12 +14,11 @@
  * limitations under the License.
  */
 
-#include "app/src/util_android.h"
-#include "firestore/src/android/firestore_android.h"
 #include "firestore/src/jni/arena_ref.h"
+#include "firestore/src/android/firestore_android.h"
 #include "firestore/src/jni/long.h"
 
-#include "firestore_integration_test_android.h"
+#include "firestore_integration_test.h"
 
 #include "gtest/gtest.h"
 
@@ -31,29 +30,29 @@ namespace {
 using ArenaRefTestAndroid = FirestoreIntegrationTest;
 
 TEST_F(ArenaRefTestAndroid, DefaultConstructorCreatesReferenceToNull) {
-  Env env = FirestoreInternal::GetEnv();
+  Env env;
   ArenaRef arena_ref;
   EXPECT_EQ(arena_ref.get(env).get(), nullptr);
 }
 
 TEST_F(ArenaRefTestAndroid, ConstructFromEnvAndObject) {
-  Env env = FirestoreInternal::GetEnv();
+  Env env;
   Local<String> string = env.NewStringUtf("hello world");
 
   ArenaRef arena_ref(env, string);
   EXPECT_TRUE(arena_ref.get(env).Equals(env, string));
 }
 
 TEST_F(ArenaRefTestAndroid, CopysReferenceToNull) {
-  Env env = FirestoreInternal::GetEnv();
+  Env env;
 
   ArenaRef arena_ref1;
   ArenaRef arena_ref2(arena_ref1);
   EXPECT_EQ(arena_ref2.get(env).get(), nullptr);
 }
 
 TEST_F(ArenaRefTestAndroid, CopysReferenceToValidObject) {
-  Env env = FirestoreInternal::GetEnv();
+  Env env;
 
   Local<String> string = env.NewStringUtf("hello world");
 
@@ -65,7 +64,7 @@ TEST_F(ArenaRefTestAndroid, CopysReferenceToValidObject) {
 }
 
 TEST_F(ArenaRefTestAndroid, CopyAssignsReferenceToNull) {
-  Env env = FirestoreInternal::GetEnv();
+  Env env;
 
   ArenaRef arena_ref1, arena_ref2;
   arena_ref2 = arena_ref1;
@@ -74,7 +73,7 @@ TEST_F(ArenaRefTestAndroid, CopyAssignsReferenceToNull) {
 }
 
 TEST_F(ArenaRefTestAndroid, CopyAssignsReferenceToValidObject) {
-  Env env = FirestoreInternal::GetEnv();
+  Env env;
 
   Local<String> string1 = env.NewStringUtf("hello world");
   Local<String> string2 = env.NewStringUtf("hello earth");
@@ -94,36 +93,33 @@ TEST_F(ArenaRefTestAndroid, CopyAssignsReferenceToValidObject) {
 }
 
 TEST_F(ArenaRefTestAndroid, MovesReferenceToNull) {
-  Env env = FirestoreInternal::GetEnv();
+  Env env;
 
   ArenaRef arena_ref1;
   ArenaRef arena_ref2(std::move(arena_ref1));
-  EXPECT_EQ(arena_ref1.get(env).get(), nullptr);
   EXPECT_EQ(arena_ref2.get(env).get(), nullptr);
 }
 
 TEST_F(ArenaRefTestAndroid, MovesReferenceToValidObject) {
-  Env env = FirestoreInternal::GetEnv();
+  Env env;
 
   Local<String> string = env.NewStringUtf("hello world");
 
   ArenaRef arena_ref2(env, string);
   ArenaRef arena_ref3(std::move(arena_ref2));
-  EXPECT_EQ(arena_ref2.get(env).get(), nullptr);
   EXPECT_TRUE(arena_ref3.get(env).Equals(env, string));
 }
 
 TEST_F(ArenaRefTestAndroid, MoveAssignsReferenceToNull) {
-  Env env = FirestoreInternal::GetEnv();
+  Env env;
 
   ArenaRef arena_ref1, arena_ref2;
   arena_ref2 = std::move(arena_ref1);
-  EXPECT_EQ(arena_ref1.get(env).get(), nullptr);
   EXPECT_EQ(arena_ref2.get(env).get(), nullptr);
 }
 
 TEST_F(ArenaRefTestAndroid, MoveAssignsReferenceToValidObject) {
-  Env env = FirestoreInternal::GetEnv();
+  Env env;
 
   Local<String> string1 = env.NewStringUtf("hello world");
   Local<String> string2 = env.NewStringUtf("hello earth");
@@ -135,7 +131,6 @@ TEST_F(ArenaRefTestAndroid, MoveAssignsReferenceToValidObject) {
 
   ArenaRef arena_ref3(env, string2);
   arena_ref3 = std::move(arena_ref2);
-  EXPECT_EQ(arena_ref2.get(env).get(), nullptr);
   EXPECT_TRUE(arena_ref3.get(env).Equals(env, string1));
 
   arena_ref3 = std::move(arena_ref1);

firestore/integration_test_internal/src/android/arena_ref_android_test.cc renamed to firestore/integration_test_internal/src/jni/arena_ref_test.cc

@@ -128,6 +128,18 @@ TEST_F(EnvTest, CallsVoidMethods) {
   EXPECT_EQ(length, 42);
 }
 
+TEST_F(EnvTest, CallsValidArenaRef) {
+  Local<String> str = env().NewStringUtf("Foo");
+  ArenaRef arena_ref(env(), str);
+
+  Local<Class> clazz = env().FindClass("java/lang/String");
+  jmethodID to_lower_case =
+      env().GetMethodId(clazz, "toLowerCase", "()Ljava/lang/String;");
+
+  Local<Object> result = env().Call(arena_ref, Method<Object>(to_lower_case));
+  EXPECT_EQ("foo", result.ToString(env()));
+}
+
 TEST_F(EnvTest, GetsStaticFields) {
   Local<Class> clazz = env().FindClass("java/lang/String");
   jfieldID comparator = env().GetStaticFieldId(clazz, "CASE_INSENSITIVE_ORDER",
@@ -160,6 +172,13 @@ TEST_F(EnvTest, ToString) {
   EXPECT_EQ("Foo", result);
 }
 
+TEST_F(EnvTest, IsInstanceOfChecksValidArenaRef) {
+  Local<Class> clazz = env().FindClass("java/lang/String");
+  Local<String> str = env().NewStringUtf("Foo");
+  ArenaRef arena_ref(env(), str);
+  EXPECT_TRUE(env().IsInstanceOf(arena_ref, clazz));
+}
+
 TEST_F(EnvTest, Throw) {
   Local<Class> clazz = env().FindClass("java/lang/Exception");
   jmethodID ctor = env().GetMethodId(clazz, "<init>", "(Ljava/lang/String;)V");

firestore/integration_test_internal/src/jni/env_test.cc

@@ -17,6 +17,7 @@
 #include "firestore/src/jni/arena_ref.h"
 
 #include <atomic>
+#include <utility>
 
 #include "firestore/src/android/firestore_android.h"
 #include "firestore/src/jni/env.h"
@@ -31,6 +32,8 @@ namespace jni {
 namespace {
 
 HashMap* gArenaRefHashMap = nullptr;
+jclass gLongClass = nullptr;
+jmethodID gLongConstructor = nullptr;
 std::mutex mutex_;
 
 int64_t GetNextArenaRefKey() {
@@ -42,80 +45,78 @@ int64_t GetNextArenaRefKey() {
 
 ArenaRef::ArenaRef(Env& env, const Object& object)
     : key_(GetNextArenaRefKey()) {
-  std::unique_lock<std::mutex> lock(mutex_);
-  gArenaRefHashMap->Put(env, key_object(env), object);
+  Local<Long> key_ref = key_object(env);
+  std::lock_guard<std::mutex> lock(mutex_);
+  gArenaRefHashMap->Put(env, key_ref, object);
 }
 
 ArenaRef::ArenaRef(const ArenaRef& other)
-    : key_(other.key_ == -1 ? -1 : GetNextArenaRefKey()) {
-  if (other.key_ != -1) {
+    : key_(other.key_ == kInvalidKey ? kInvalidKey : GetNextArenaRefKey()) {
+  if (other.key_ != kInvalidKey) {
     Env env = FirestoreInternal::GetEnv();
     Local<Object> object = other.get(env);
-    std::unique_lock<std::mutex> lock(mutex_);
-    gArenaRefHashMap->Put(env, key_object(env), object);
+    Local<Long> key_ref = key_object(env);
+    std::lock_guard<std::mutex> lock(mutex_);
+    gArenaRefHashMap->Put(env, key_ref, object);
   }
 }
 
 ArenaRef::ArenaRef(ArenaRef&& other) { std::swap(key_, other.key_); }
 
 ArenaRef& ArenaRef::operator=(const ArenaRef& other) {
-  Env env = FirestoreInternal::GetEnv();
-
-  if (this == &other) {
+  if (this == &other || (key_ == kInvalidKey && other.key_ == kInvalidKey)) {
     return *this;
   }
 
-  {
-    std::unique_lock<std::mutex> lock(mutex_);
-    if (key_ != -1) {
-      gArenaRefHashMap->Remove(env, key_object(env));
-      key_ = -1;
-    }
-
-    if (other.key_ != -1) {
+  Env env = FirestoreInternal::GetEnv();
+  if (other.key_ != kInvalidKey) {
+    if (key_ == kInvalidKey) {
       key_ = GetNextArenaRefKey();
-      Local<Object> object = other.get(env);
-      gArenaRefHashMap->Put(env, key_object(env), object);
     }
+    Local<Object> object = other.get(env);
+    std::lock_guard<std::mutex> lock(mutex_);
+    gArenaRefHashMap->Put(env, key_object(env), object);
+  } else {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      gArenaRefHashMap->Remove(env, key_object(env));
+    }
+    key_ = kInvalidKey;
   }
   return *this;
 }
 
 ArenaRef& ArenaRef::operator=(ArenaRef&& other) {
-  Env env = FirestoreInternal::GetEnv();
-
-  if (this == &other) {
-    return *this;
-  }
-
-  {
-    std::unique_lock<std::mutex> lock(mutex_);
-    if (key_ != -1) {
-      gArenaRefHashMap->Remove(env, key_object(env));
-    }
-    key_ = other.key_;
-    other.key_ = -1;
+  if (this != &other) {
+    std::swap(key_, other.key_);
   }
   return *this;
 }
 
 ArenaRef::~ArenaRef() {
-  if (key_ != -1) {
+  if (key_ != kInvalidKey) {
     Env env;
     ExceptionClearGuard block(env);
     {
-      std::unique_lock<std::mutex> lock(mutex_);
+      std::lock_guard<std::mutex> lock(mutex_);
       gArenaRefHashMap->Remove(env, key_object(env));
     }
     if (!env.ok()) {
-      env.ExceptionDescribe();
+      env.RecordException();
       env.ExceptionClear();
     }
   }
 }
 
 Local<Long> ArenaRef::key_object(Env& env) const {
-  return Long::Create(env, key_);
+  if (!env.ok()) return {};
+  // The reason why using raw JNI calls here rather than the JNI convenience
+  // layer is because we need to get rid of JNI convenience layer dependencies
+  // in the ArenaRef destructor to avoid racing condition when firestore object
+  // gets destroyed.
+  jobject key = env.get()->NewObject(gLongClass, gLongConstructor, key_);
+  env.RecordException();
+  return {env.get(), key};
 }
 
 void ArenaRef::Initialize(Env& env) {
@@ -124,6 +125,13 @@ void ArenaRef::Initialize(Env& env) {
   }
   Global<HashMap> hash_map(HashMap::Create(env));
   gArenaRefHashMap = new HashMap(hash_map.release());
+
+  auto long_class_local = env.get()->FindClass("java/lang/Long");
+  gLongClass = static_cast<jclass>(env.get()->NewGlobalRef(long_class_local));
+  env.get()->DeleteLocalRef(long_class_local);
+  gLongConstructor = env.get()->GetMethodID(gLongClass, "<init>", "(J)V");
+
+  if (!env.ok()) return;
 }
 
 Local<Object> ArenaRef::get(Env& env) const {