Implement ObjectArena and ArenaRef to solve JNI Global Reference Issue

firestore/CMakeLists.txt

@@ -138,6 +138,8 @@ set(android_SRCS
     src/android/wrapper.h
     src/android/write_batch_android.cc
     src/android/write_batch_android.h
+    src/jni/arena_ref.cc
+    src/jni/arena_ref.h
     src/jni/array.h
     src/jni/array_list.cc
     src/jni/array_list.h

firestore/integration_test_internal/CMakeLists.txt

@@ -131,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/document_reference_test.cc

@@ -59,7 +59,7 @@ TEST_F(DocumentReferenceTest, RecoverFirestore) {
   DocumentReference doc = Document();
   ASSERT_EQ(db, doc.firestore());  // Sanity check
 
-  jni::Object doc_java = GetInternal(doc)->ToJava();
+  jni::Local<jni::Object> doc_java = GetInternal(doc)->ToJava();
   result = DocumentReferenceInternal::Create(env, doc_java);
   ASSERT_EQ(db, result.firestore());
 }

firestore/integration_test_internal/src/jni/arena_ref_test.cc

@@ -0,0 +1,143 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "firestore/src/jni/arena_ref.h"
+#include "firestore/src/android/firestore_android.h"
+#include "firestore/src/jni/long.h"
+
+#include "firestore_integration_test.h"
+
+#include "gtest/gtest.h"
+
+namespace firebase {
+namespace firestore {
+namespace jni {
+namespace {
+
+using ArenaRefTestAndroid = FirestoreIntegrationTest;
+
+TEST_F(ArenaRefTestAndroid, DefaultConstructorCreatesReferenceToNull) {
+  Env env;
+  ArenaRef arena_ref;
+  EXPECT_EQ(arena_ref.get(env).get(), nullptr);
+}
+
+TEST_F(ArenaRefTestAndroid, ConstructFromEnvAndObject) {
+  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;
+
+  ArenaRef arena_ref1;
+  ArenaRef arena_ref2(arena_ref1);
+  EXPECT_EQ(arena_ref2.get(env).get(), nullptr);
+}
+
+TEST_F(ArenaRefTestAndroid, CopysReferenceToValidObject) {
+  Env env;
+
+  Local<String> string = env.NewStringUtf("hello world");
+
+  ArenaRef arena_ref1(env, string);
+  ArenaRef arena_ref2(arena_ref1);
+
+  EXPECT_TRUE(arena_ref1.get(env).Equals(env, string));
+  EXPECT_TRUE(arena_ref2.get(env).Equals(env, string));
+}
+
+TEST_F(ArenaRefTestAndroid, CopyAssignsReferenceToNull) {
+  Env env;
+
+  ArenaRef arena_ref1, arena_ref2;
+  arena_ref2 = arena_ref1;
+  EXPECT_EQ(arena_ref1.get(env).get(), nullptr);
+  EXPECT_EQ(arena_ref2.get(env).get(), nullptr);
+}
+
+TEST_F(ArenaRefTestAndroid, CopyAssignsReferenceToValidObject) {
+  Env env;
+
+  Local<String> string1 = env.NewStringUtf("hello world");
+  Local<String> string2 = env.NewStringUtf("hello earth");
+
+  ArenaRef arena_ref1;
+  ArenaRef arena_ref2(env, string1);
+  ArenaRef arena_ref3(env, string2);
+  arena_ref3 = arena_ref2;
+  arena_ref2 = arena_ref2;
+
+  EXPECT_TRUE(arena_ref3.get(env).Equals(env, string1));
+  EXPECT_TRUE(arena_ref2.get(env).Equals(env, string1));
+
+  arena_ref2 = arena_ref1;
+  EXPECT_EQ(arena_ref2.get(env).get(), nullptr);
+  EXPECT_TRUE(arena_ref3.get(env).Equals(env, string1));
+}
+
+TEST_F(ArenaRefTestAndroid, MovesReferenceToNull) {
+  Env env;
+
+  ArenaRef arena_ref1;
+  ArenaRef arena_ref2(std::move(arena_ref1));
+  EXPECT_EQ(arena_ref2.get(env).get(), nullptr);
+}
+
+TEST_F(ArenaRefTestAndroid, MovesReferenceToValidObject) {
+  Env env;
+
+  Local<String> string = env.NewStringUtf("hello world");
+
+  ArenaRef arena_ref2(env, string);
+  ArenaRef arena_ref3(std::move(arena_ref2));
+  EXPECT_TRUE(arena_ref3.get(env).Equals(env, string));
+}
+
+TEST_F(ArenaRefTestAndroid, MoveAssignsReferenceToNull) {
+  Env env;
+
+  ArenaRef arena_ref1, arena_ref2;
+  arena_ref2 = std::move(arena_ref1);
+  EXPECT_EQ(arena_ref2.get(env).get(), nullptr);
+}
+
+TEST_F(ArenaRefTestAndroid, MoveAssignsReferenceToValidObject) {
+  Env env;
+
+  Local<String> string1 = env.NewStringUtf("hello world");
+  Local<String> string2 = env.NewStringUtf("hello earth");
+
+  ArenaRef arena_ref1;
+  ArenaRef arena_ref2(env, string1);
+  arena_ref2 = std::move(arena_ref2);
+  EXPECT_TRUE(arena_ref2.get(env).Equals(env, string1));
+
+  ArenaRef arena_ref3(env, string2);
+  arena_ref3 = std::move(arena_ref2);
+  EXPECT_TRUE(arena_ref3.get(env).Equals(env, string1));
+
+  arena_ref3 = std::move(arena_ref1);
+  EXPECT_EQ(arena_ref3.get(env).get(), nullptr);
+}
+
+}  // namespace
+}  // namespace jni
+}  // namespace firestore
+}  // namespace firebase

firestore/integration_test_internal/src/jni/env_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/src/android/document_reference_android.cc

@@ -147,14 +147,14 @@ Future<void> DocumentReferenceInternal::Set(const MapFieldValue& data,
   Env env = GetEnv();
   FieldValueInternal map_value(data);
   Local<Object> java_options = SetOptionsInternal::Create(env, options);
-  Local<Task> task = env.Call(obj_, kSet, map_value, java_options);
+  Local<Task> task = env.Call(obj_, kSet, map_value.ToJava(), java_options);
   return promises_.NewFuture<void>(env, AsyncFn::kSet, task);
 }
 
 Future<void> DocumentReferenceInternal::Update(const MapFieldValue& data) {
   Env env = GetEnv();
   FieldValueInternal map_value(data);
-  Local<Task> task = env.Call(obj_, kUpdate, map_value);
+  Local<Task> task = env.Call(obj_, kUpdate, map_value.ToJava());
   return promises_.NewFuture<void>(env, AsyncFn::kUpdate, task);
 }
 

firestore/src/android/field_value_android.cc

@@ -38,6 +38,7 @@ namespace firebase {
 namespace firestore {
 namespace {
 
+using jni::ArenaRef;
 using jni::Array;
 using jni::ArrayList;
 using jni::Boolean;
@@ -98,39 +99,45 @@ FieldValue FieldValueInternal::Create(Env& env,
 FieldValueInternal::FieldValueInternal() : cached_type_(Type::kNull) {}
 
 FieldValueInternal::FieldValueInternal(const Object& object)
-    : object_(object), cached_type_(Type::kNull) {}
+    : cached_type_(Type::kNull) {
+  Env env = GetEnv();
+  object_ = ArenaRef(env, object);
+}
 
 FieldValueInternal::FieldValueInternal(Type type, const Object& object)
-    : object_(object), cached_type_(type) {}
+    : cached_type_(type) {
+  Env env = GetEnv();
+  object_ = ArenaRef(env, object);
+}
 
 FieldValueInternal::FieldValueInternal(bool value)
     : cached_type_(Type::kBoolean) {
   Env env = GetEnv();
-  object_ = Boolean::Create(env, value);
+  object_ = ArenaRef(env, Boolean::Create(env, value));
 }
 
 FieldValueInternal::FieldValueInternal(int64_t value)
     : cached_type_(Type::kInteger) {
   Env env = GetEnv();
-  object_ = Long::Create(env, value);
+  object_ = ArenaRef(env, Long::Create(env, value));
 }
 
 FieldValueInternal::FieldValueInternal(double value)
     : cached_type_(Type::kDouble) {
   Env env = GetEnv();
-  object_ = Double::Create(env, value);
+  object_ = ArenaRef(env, Double::Create(env, value));
 }
 
 FieldValueInternal::FieldValueInternal(Timestamp value)
     : cached_type_(Type::kTimestamp) {
   Env env = GetEnv();
-  object_ = TimestampInternal::Create(env, value);
+  object_ = ArenaRef(env, TimestampInternal::Create(env, value));
 }
 
 FieldValueInternal::FieldValueInternal(std::string value)
     : cached_type_(Type::kString) {
   Env env = GetEnv();
-  object_ = env.NewStringUtf(value);
+  object_ = ArenaRef(env, env.NewStringUtf(value));
 }
 
 // We do not initialize cached_blob_ with value here as the instance constructed
@@ -140,20 +147,21 @@ FieldValueInternal::FieldValueInternal(std::string value)
 FieldValueInternal::FieldValueInternal(const uint8_t* value, size_t size)
     : cached_type_(Type::kBlob) {
   Env env = GetEnv();
-  object_ = BlobInternal::Create(env, value, size);
+  object_ = ArenaRef(env, BlobInternal::Create(env, value, size));
 }
 
 FieldValueInternal::FieldValueInternal(DocumentReference value)
     : cached_type_{Type::kReference} {
   if (value.internal_ != nullptr) {
-    object_ = value.internal_->ToJava();
+    Env env = GetEnv();
+    object_ = ArenaRef(env, value.internal_->ToJava());
   }
 }
 
 FieldValueInternal::FieldValueInternal(GeoPoint value)
     : cached_type_(Type::kGeoPoint) {
   Env env = GetEnv();
-  object_ = GeoPointInternal::Create(env, value);
+  object_ = ArenaRef(env, GeoPointInternal::Create(env, value));
 }
 
 FieldValueInternal::FieldValueInternal(const std::vector<FieldValue>& value)
@@ -164,7 +172,7 @@ FieldValueInternal::FieldValueInternal(const std::vector<FieldValue>& value)
     // TODO(b/150016438): don't conflate invalid `FieldValue`s and null.
     list.Add(env, ToJava(element));
   }
-  object_ = list;
+  object_ = ArenaRef(env, list);
 }
 
 FieldValueInternal::FieldValueInternal(const MapFieldValue& value)
@@ -176,20 +184,21 @@ FieldValueInternal::FieldValueInternal(const MapFieldValue& value)
     Local<String> key = env.NewStringUtf(kv.first);
     map.Put(env, key, ToJava(kv.second));
   }
-  object_ = map;
+  object_ = ArenaRef(env, map);
 }
 
 Type FieldValueInternal::type() const {
   if (cached_type_ != Type::kNull) {
     return cached_type_;
   }
-  if (!object_) {
+
+  Env env = GetEnv();
+  if (!object_.get(env)) {
     return Type::kNull;
   }
 
   // We do not have any knowledge on the type yet. Check the runtime type with
   // each known type.
-  Env env = GetEnv();
   if (env.IsInstanceOf(object_, Boolean::GetClass())) {
     cached_type_ = Type::kBoolean;
     return Type::kBoolean;
@@ -232,7 +241,7 @@ Type FieldValueInternal::type() const {
   }
 
   FIREBASE_ASSERT_MESSAGE(false, "Unsupported FieldValue type: %s",
-                          Class::GetClassName(env, object_).c_str());
+                          Class::GetClassName(env, object_.get(env)).c_str());
   return Type::kNull;
 }
 
@@ -389,15 +398,25 @@ bool operator==(const FieldValueInternal& lhs, const FieldValueInternal& rhs) {
 }
 
 template <typename T>
-T FieldValueInternal::Cast(jni::Env& env, Type type) const {
+Local<T> FieldValueInternal::Cast(jni::Env& env, Type type) const {
   if (cached_type_ == Type::kNull) {
     FIREBASE_ASSERT(env.IsInstanceOf(object_, T::GetClass()));
     cached_type_ = type;
   } else {
     FIREBASE_ASSERT(cached_type_ == type);
   }
-  auto typed_value = static_cast<jni::JniType<T>>(object_.get());
-  return T(typed_value);
+  return object_.get(env).CastTo<T>();
+}
+
+template <>
+Local<String> FieldValueInternal::Cast<String>(jni::Env& env, Type type) const {
+  if (cached_type_ == Type::kNull) {
+    FIREBASE_ASSERT(env.IsInstanceOf(object_, String::GetClass()));
+    cached_type_ = type;
+  } else {
+    FIREBASE_ASSERT(cached_type_ == type);
+  }
+  return env.NewStringUtf(object_.get(env).ToString(env));
 }
 
 Local<Array<Object>> FieldValueInternal::MakeArray(
@@ -411,8 +430,14 @@ Local<Array<Object>> FieldValueInternal::MakeArray(
 
 Env FieldValueInternal::GetEnv() { return FirestoreInternal::GetEnv(); }
 
-Object FieldValueInternal::ToJava(const FieldValue& value) {
-  return value.internal_ ? value.internal_->object_ : Object();
+Local<Object> FieldValueInternal::ToJava() const {
+  Env env = GetEnv();
+  return object_.get(env);
+}
+
+Local<Object> FieldValueInternal::ToJava(const FieldValue& value) {
+  Env env = GetEnv();
+  return value.internal_ ? value.internal_->object_.get(env) : Local<Object>();
 }
 
 }  // namespace firestore