[remoteobjects] Add a reference count's map in RemoteObjectHostImpl

This is precursor CL of https://crrev.com/c/2379270, and adds the
reference count's map to remove the object from the registry when
the reference count is zero.

Bug: 1105925
Change-Id: I5b04c1fe17c4376dad97dedf3f5750e05f611c32
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2417740
Commit-Queue: Miyoung Shin <myid.shin@igalia.com>
Reviewed-by: Jeremy Roman <jbroman@chromium.org>
Reviewed-by: Shimi Zhang <ctzsm@chromium.org>
Reviewed-by: Kinuko Yasuda <kinuko@chromium.org>
Reviewed-by: Oksana Zhuravlova <oksamyt@chromium.org>
Cr-Commit-Position: refs/heads/master@{#814521}

content/browser/renderer_host/render_frame_host_impl_browsertest.cc

@@ -4768,10 +4768,15 @@ class MockObjectHost : public blink::mojom::RemoteObjectHost {
       mojo::MakeSelfOwnedReceiver(std::make_unique<MockInnerObject>(),
                                   std::move(receiver));
     }
+    reference_count_map_[object_id]++;
+  }
+
+  void AcquireObject(int32_t object_id) override {
+    reference_count_map_[object_id]++;
   }
 
   void ReleaseObject(int32_t object_id) override {
-    release_object_called_[object_id] = true;
+    reference_count_map_[object_id]--;
   }
 
   mojo::PendingRemote<blink::mojom::RemoteObjectHost> GetRemote() {
@@ -4780,15 +4785,15 @@ class MockObjectHost : public blink::mojom::RemoteObjectHost {
 
   MockObject* GetMockObject() const { return mock_object_.get(); }
 
-  bool release_object_called_for_object(int32_t object_id) const {
-    return release_object_called_.at(object_id);
+  int ReferenceCount(int32_t object_id) const {
+    return !reference_count_map_.at(object_id);
   }
 
  private:
   mojo::Receiver<blink::mojom::RemoteObjectHost> receiver_{this};
   std::unique_ptr<MockObject> mock_object_;
-  std::map<int32_t, bool> release_object_called_{{kMainObject.id, false},
-                                                 {kInnerObject.id, false}};
+  std::map<int32_t, int> reference_count_map_{{kMainObject.id, 0},
+                                              {kInnerObject.id, 0}};
 };
 
 class RemoteObjectInjector : public WebContentsObserver {
@@ -4926,16 +4931,13 @@ IN_PROC_BROWSER_TEST_F(RenderFrameHostImplBrowserTest, RemoteObjectRelease) {
           web_contents(),
           "globalInner = testObject.getInnerObject(); typeof globalInner; "));
 
-  EXPECT_FALSE(injector.GetObjectHost().release_object_called_for_object(
-      kInnerObject.id));
+  EXPECT_GT(injector.GetObjectHost().ReferenceCount(kInnerObject.id), 0);
   EXPECT_EQ("object", EvalJs(web_contents(), "gc(); typeof globalInner;"));
-  EXPECT_FALSE(injector.GetObjectHost().release_object_called_for_object(
-      kInnerObject.id));
+  EXPECT_GT(injector.GetObjectHost().ReferenceCount(kInnerObject.id), 0);
   EXPECT_EQ(
       "undefined",
       EvalJs(web_contents(), "delete globalInner; gc(); typeof globalInner;"));
-  EXPECT_TRUE(injector.GetObjectHost().release_object_called_for_object(
-      kInnerObject.id));
+  EXPECT_EQ(injector.GetObjectHost().ReferenceCount(kInnerObject.id), 0);
 }
 
 #endif  // OS_ANDROID

content/public/android/java/src/org/chromium/content/browser/remoteobjects/RemoteObjectHostImpl.java

@@ -79,13 +79,22 @@ class RemoteObjectHostImpl implements RemoteObjectHost {
         }
     }
 
+    @Override
+    public void acquireObject(int objectId) {
+        RemoteObjectRegistry registry = mRegistry.get();
+        if (registry == null) {
+            return;
+        }
+        registry.refObjectById(objectId);
+    }
+
     @Override
     public void releaseObject(int objectId) {
         RemoteObjectRegistry registry = mRegistry.get();
         if (registry == null) {
             return;
         }
-        registry.removeObjectById(objectId);
+        registry.unrefObjectById(objectId);
     }
 
     @Override

content/public/android/java/src/org/chromium/content/browser/remoteobjects/RemoteObjectRegistry.java

@@ -4,7 +4,9 @@
 
 package org.chromium.content.browser.remoteobjects;
 
-import java.util.HashMap;
+import android.util.SparseArray;
+
+import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.Set;
 
@@ -16,8 +18,29 @@ import java.util.Set;
  */
 final class RemoteObjectRegistry implements RemoteObjectImpl.ObjectIdAllocator {
     private final Set<? super RemoteObjectRegistry> mRetainingSet;
-    private final Map<Integer, Object> mObjectsById = new HashMap<>();
-    private final Map<Object, Integer> mIdsByObject = new HashMap<>();
+
+    private static class Entry {
+        Entry(int id, Object object) {
+            this.id = id;
+            this.object = object;
+        }
+
+        // The ID assigned to the object, which can be used by the renderer
+        // to refer to it.
+        public final int id;
+
+        // The injected object itself
+        public final Object object;
+
+        // The number of outstanding references to this object.
+        // This includes:
+        //   * wrapper objects in the renderer (removed when it closes the pipe)
+        //   * names assigned to the object by developer Java code
+        public int referenceCount = 1;
+    }
+
+    private final SparseArray<Entry> mEntriesById = new SparseArray<>();
+    private final Map<Object, Entry> mEntriesByObject = new IdentityHashMap<>();
     private int mNextId;
 
     RemoteObjectRegistry(Set<? super RemoteObjectRegistry> retainingSet) {
@@ -31,26 +54,43 @@ final class RemoteObjectRegistry implements RemoteObjectImpl.ObjectIdAllocator {
     }
 
     @Override
-    public int getObjectId(Object object) {
-        Integer existingId = mIdsByObject.get(object);
-        if (existingId != null) {
-            return existingId;
+    public synchronized int getObjectId(Object object) {
+        Entry entry = mEntriesByObject.get(object);
+        if (entry != null) {
+            entry.referenceCount++;
+            return entry.id;
         }
 
         int newId = mNextId++;
         assert newId >= 0;
-        mObjectsById.put(newId, object);
-        mIdsByObject.put(object, newId);
+        entry = new Entry(newId, object);
+        mEntriesById.put(newId, entry);
+        mEntriesByObject.put(object, entry);
         return newId;
     }
 
     @Override
-    public Object getObjectById(int id) {
-        return mObjectsById.get(id);
+    public synchronized Object getObjectById(int id) {
+        Entry entry = mEntriesById.get(id);
+        return entry != null ? entry.object : null;
+    }
+
+    public synchronized void refObjectById(int id) {
+        Entry entry = mEntriesById.get(id);
+        if (entry == null) return;
+
+        entry.referenceCount++;
     }
 
-    public void removeObjectById(int id) {
-        Object o = mObjectsById.remove(id);
-        mIdsByObject.remove(o);
+    public synchronized void unrefObjectById(int id) {
+        Entry entry = mEntriesById.get(id);
+        if (entry == null) return;
+
+        entry.referenceCount--;
+        assert entry.referenceCount >= 0;
+        if (entry.referenceCount > 0) return;
+
+        mEntriesById.remove(entry.id);
+        mEntriesByObject.remove(entry.object);
     }
 }

content/public/android/junit/src/org/chromium/content/browser/remoteobjects/RemoteObjectRegistryTest.java

@@ -10,7 +10,8 @@ import static org.hamcrest.Matchers.not;
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.BlockJUnit4ClassRunner;
+
+import org.chromium.base.test.BaseRobolectricTestRunner;
 
 import java.util.HashSet;
 import java.util.Set;
@@ -18,7 +19,7 @@ import java.util.Set;
 /**
  * Tests the object registry, which maintains bidirectional object/ID mappings.
  */
-@RunWith(BlockJUnit4ClassRunner.class)
+@RunWith(BaseRobolectricTestRunner.class)
 public final class RemoteObjectRegistryTest {
     @Test
     public void testMaintainsRetainingSet() {
@@ -58,7 +59,7 @@ public final class RemoteObjectRegistryTest {
         RemoteObjectRegistry registry = new RemoteObjectRegistry(retainingSet);
         Object o = new Object();
         int id = registry.getObjectId(o);
-        registry.removeObjectById(id);
+        registry.unrefObjectById(id);
         int id2 = registry.getObjectId(o);
         Assert.assertSame(o, registry.getObjectById(id2));
     }

third_party/blink/public/mojom/remote_objects/remote_objects.mojom

@@ -21,6 +21,9 @@ interface RemoteObjectHost {
   // |object_id|.
   GetObject(int32 object_id, pending_receiver<RemoteObject> receiver);
 
+  // Increase the object's reference count when the named object is injected to JS.
+  AcquireObject(int32 object_id);
+
   // Ideally lifetime would be scoped to an associated RemoteObject pipe. This
   // is complicated by the fact that the missing support for associated
   // interfaces in the Mojo Java bindings.

third_party/blink/renderer/modules/remote_objects/remote_object_gateway_impl.cc

@@ -45,6 +45,7 @@ void RemoteObjectGatewayImpl::InjectNamed(const WTF::String& object_name,
 
   global->Set(context, V8AtomicString(isolate, object_name), controller.ToV8())
       .Check();
+  object_host_->AcquireObject(object_id);
 }
 
 // static