[vm/concurrency] Make updates to MegamorphicCache use type_feedback_mutex to prevent concurrent access

With lightweights isolates enabled we need to ensure to update entries
in MetamorphicCache in sequence with mutual exclusion.

This CL adds appropriate locking for the MegamorphicCache.

Right now this is safe in AOT because the entire SwitchableCallMiss
handler runs with stopped mutators. A future CL will relax this. That
future CL will also move the megamorphic mutex from Isolate to
IsolateGroup, since the megamoprhic caches are shared within an isolate
group.

Issue https://github.com/dart-lang/sdk/issues/36097

TEST=CL is purely a refactoring.

Change-Id: I30305b04027d2bf9b4d5f6b1700bf158d38bdba7
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/174467
Commit-Queue: Martin Kustermann <kustermann@google.com>
Reviewed-by: Vyacheslav Egorov <vegorov@google.com>
diff --git a/runtime/vm/compiler/backend/il.cc b/runtime/vm/compiler/backend/il.cc
index bd8e8b7..9fdefe6 100644
--- a/runtime/vm/compiler/backend/il.cc
+++ b/runtime/vm/compiler/backend/il.cc
@@ -820,23 +820,26 @@
     const Array& descriptor =
         Array::Handle(zone, ic_data.arguments_descriptor());
     Thread* thread = Thread::Current();
-    const MegamorphicCache& cache = MegamorphicCache::Handle(
+
+    const auto& cache = MegamorphicCache::Handle(
         zone, MegamorphicCacheTable::Lookup(thread, name, descriptor));
-    SafepointMutexLocker ml(thread->isolate()->megamorphic_mutex());
-    MegamorphicCacheEntries entries(Array::Handle(zone, cache.buckets()));
-    for (intptr_t i = 0, n = entries.Length(); i < n; i++) {
-      const intptr_t id =
-          Smi::Value(entries[i].Get<MegamorphicCache::kClassIdIndex>());
-      if (id == kIllegalCid) {
-        continue;
+    {
+      SafepointMutexLocker ml(thread->isolate()->type_feedback_mutex());
+      MegamorphicCacheEntries entries(Array::Handle(zone, cache.buckets()));
+      for (intptr_t i = 0, n = entries.Length(); i < n; i++) {
+        const intptr_t id =
+            Smi::Value(entries[i].Get<MegamorphicCache::kClassIdIndex>());
+        if (id == kIllegalCid) {
+          continue;
+        }
+        Function& function = Function::ZoneHandle(zone);
+        function ^= entries[i].Get<MegamorphicCache::kTargetFunctionIndex>();
+        const intptr_t filled_entry_count = cache.filled_entry_count();
+        ASSERT(filled_entry_count > 0);
+        cid_ranges_.Add(new (zone) TargetInfo(
+            id, id, &function, Usage(function) / filled_entry_count,
+            StaticTypeExactnessState::NotTracking()));
       }
-      Function& function = Function::ZoneHandle(zone);
-      function ^= entries[i].Get<MegamorphicCache::kTargetFunctionIndex>();
-      const intptr_t filled_entry_count = cache.filled_entry_count();
-      ASSERT(filled_entry_count > 0);
-      cid_ranges_.Add(new (zone) TargetInfo(
-          id, id, &function, Usage(function) / filled_entry_count,
-          StaticTypeExactnessState::NotTracking()));
     }
   }
 }
diff --git a/runtime/vm/isolate.cc b/runtime/vm/isolate.cc
index f8fa4cb..3849287 100644
--- a/runtime/vm/isolate.cc
+++ b/runtime/vm/isolate.cc
@@ -1648,7 +1648,8 @@
       mutex_(NOT_IN_PRODUCT("Isolate::mutex_")),
       constant_canonicalization_mutex_(
           NOT_IN_PRODUCT("Isolate::constant_canonicalization_mutex_")),
-      megamorphic_mutex_(NOT_IN_PRODUCT("Isolate::megamorphic_mutex_")),
+      megamorphic_table_mutex_(
+          NOT_IN_PRODUCT("Isolate::megamorphic_table_mutex_")),
       type_feedback_mutex_(NOT_IN_PRODUCT("Isolate::type_feedback_mutex_")),
       kernel_data_lib_cache_mutex_(
           NOT_IN_PRODUCT("Isolate::kernel_data_lib_cache_mutex_")),
diff --git a/runtime/vm/isolate.h b/runtime/vm/isolate.h
index 93eb508..77991a9 100644
--- a/runtime/vm/isolate.h
+++ b/runtime/vm/isolate.h
@@ -964,7 +964,7 @@
   Mutex* constant_canonicalization_mutex() {
     return &constant_canonicalization_mutex_;
   }
-  Mutex* megamorphic_mutex() { return &megamorphic_mutex_; }
+  Mutex* megamorphic_table_mutex() { return &megamorphic_table_mutex_; }
   Mutex* type_feedback_mutex() { return &type_feedback_mutex_; }
 
   Mutex* kernel_data_lib_cache_mutex() { return &kernel_data_lib_cache_mutex_; }
@@ -1601,7 +1601,7 @@
   Simulator* simulator_ = nullptr;
   Mutex mutex_;                            // Protects compiler stats.
   Mutex constant_canonicalization_mutex_;  // Protects const canonicalization.
-  Mutex megamorphic_mutex_;
+  Mutex megamorphic_table_mutex_;
   Mutex type_feedback_mutex_;
   Mutex kernel_data_lib_cache_mutex_;
   Mutex kernel_data_class_cache_mutex_;
diff --git a/runtime/vm/megamorphic_cache_table.cc b/runtime/vm/megamorphic_cache_table.cc
index 04a4acd..7a1b7f6 100644
--- a/runtime/vm/megamorphic_cache_table.cc
+++ b/runtime/vm/megamorphic_cache_table.cc
@@ -17,13 +17,13 @@
                                                   const String& name,
                                                   const Array& descriptor) {
   Isolate* isolate = thread->isolate();
-  // Multiple compilation threads could access this lookup.
-  SafepointMutexLocker ml(isolate->megamorphic_mutex());
+  SafepointMutexLocker ml(isolate->megamorphic_table_mutex());
+
   ASSERT(name.IsSymbol());
   // TODO(rmacnak): ASSERT(descriptor.IsCanonical());
 
   // TODO(rmacnak): Make a proper hashtable a la symbol table.
-  GrowableObjectArray& table = GrowableObjectArray::Handle(
+  auto& table = GrowableObjectArray::Handle(
       isolate->object_store()->megamorphic_cache_table());
   MegamorphicCache& cache = MegamorphicCache::Handle();
   if (table.IsNull()) {
@@ -45,7 +45,10 @@
 }
 
 void MegamorphicCacheTable::PrintSizes(Isolate* isolate) {
-  StackZone zone(Thread::Current());
+  auto thread = Thread::Current();
+  SafepointMutexLocker ml(thread->isolate()->megamorphic_table_mutex());
+
+  StackZone zone(thread);
   intptr_t size = 0;
   MegamorphicCache& cache = MegamorphicCache::Handle();
   Array& buckets = Array::Handle();
diff --git a/runtime/vm/object.cc b/runtime/vm/object.cc
index 1257675..3f6ff74 100644
--- a/runtime/vm/object.cc
+++ b/runtime/vm/object.cc
@@ -17172,14 +17172,78 @@
   return result.raw();
 }
 
-void MegamorphicCache::Insert(const Smi& class_id, const Object& target) const {
-  SafepointMutexLocker ml(Isolate::Current()->megamorphic_mutex());
-  EnsureCapacityLocked();
-  InsertLocked(class_id, target);
+void MegamorphicCache::EnsureContains(const Smi& class_id,
+                                      const Object& target) const {
+  SafepointMutexLocker ml(Isolate::Current()->type_feedback_mutex());
+
+  if (LookupLocked(class_id) == Object::null()) {
+    InsertLocked(class_id, target);
+  }
+
+#if defined(DEBUG)
+  if (FLAG_precompiled_mode && FLAG_use_bare_instructions) {
+    if (target.IsFunction()) {
+      const auto& function = Function::Cast(target);
+      const auto& entry_point = Smi::Handle(
+          Smi::FromAlignedAddress(Code::EntryPointOf(function.CurrentCode())));
+      ASSERT(LookupLocked(class_id) == entry_point.raw());
+    }
+  } else {
+    ASSERT(LookupLocked(class_id) == target.raw());
+  }
+#endif  // define(DEBUG)
+}
+
+ObjectPtr MegamorphicCache::Lookup(const Smi& class_id) const {
+  SafepointMutexLocker ml(Isolate::Current()->type_feedback_mutex());
+  return LookupLocked(class_id);
+}
+
+ObjectPtr MegamorphicCache::LookupLocked(const Smi& class_id) const {
+  auto thread = Thread::Current();
+  auto zone = thread->zone();
+  ASSERT(thread->IsMutatorThread());
+  ASSERT(thread->isolate()->type_feedback_mutex()->IsOwnedByCurrentThread());
+
+  const auto& backing_array = Array::Handle(zone, buckets());
+  intptr_t id_mask = mask();
+  intptr_t index = (class_id.Value() * kSpreadFactor) & id_mask;
+  intptr_t i = index;
+  do {
+    const classid_t current_cid =
+        Smi::Value(Smi::RawCast(GetClassId(backing_array, i)));
+    if (current_cid == class_id.Value()) {
+      return GetTargetFunction(backing_array, i);
+    } else if (current_cid == kIllegalCid) {
+      return Object::null();
+    }
+    i = (i + 1) & id_mask;
+  } while (i != index);
+  UNREACHABLE();
+}
+
+void MegamorphicCache::InsertLocked(const Smi& class_id,
+                                    const Object& target) const {
+  auto isolate_group = IsolateGroup::Current();
+  ASSERT(Isolate::Current()->type_feedback_mutex()->IsOwnedByCurrentThread());
+
+  // As opposed to ICData we are stopping mutator threads from other isolates
+  // while modifying the megamorphic cache, since updates are not atomic.
+  //
+  // NOTE: In the future we might change the megamorphic cache insertions to
+  // carefully use store-release barriers on the writer as well as
+  // load-acquire barriers on the reader, ...
+  isolate_group->RunWithStoppedMutators(
+      [&]() {
+        EnsureCapacityLocked();
+        InsertEntryLocked(class_id, target);
+      },
+      /*use_force_growth=*/true);
 }
 
 void MegamorphicCache::EnsureCapacityLocked() const {
-  ASSERT(Isolate::Current()->megamorphic_mutex()->IsOwnedByCurrentThread());
+  ASSERT(Isolate::Current()->type_feedback_mutex()->IsOwnedByCurrentThread());
+
   intptr_t old_capacity = mask() + 1;
   double load_limit = kLoadFactor * static_cast<double>(old_capacity);
   if (static_cast<double>(filled_entry_count() + 1) > load_limit) {
@@ -17202,15 +17266,16 @@
       class_id ^= GetClassId(old_buckets, i);
       if (class_id.Value() != kIllegalCid) {
         target = GetTargetFunction(old_buckets, i);
-        InsertLocked(class_id, target);
+        InsertEntryLocked(class_id, target);
       }
     }
   }
 }
 
-void MegamorphicCache::InsertLocked(const Smi& class_id,
-                                    const Object& target) const {
-  ASSERT(Isolate::Current()->megamorphic_mutex()->IsOwnedByCurrentThread());
+void MegamorphicCache::InsertEntryLocked(const Smi& class_id,
+                                         const Object& target) const {
+  ASSERT(Isolate::Current()->type_feedback_mutex()->IsOwnedByCurrentThread());
+
   ASSERT(Thread::Current()->IsMutatorThread());
   ASSERT(static_cast<double>(filled_entry_count() + 1) <=
          (kLoadFactor * static_cast<double>(mask() + 1)));
diff --git a/runtime/vm/object.h b/runtime/vm/object.h
index 5ca1a8c..ada80f9 100644
--- a/runtime/vm/object.h
+++ b/runtime/vm/object.h
@@ -6850,7 +6850,8 @@
   static MegamorphicCachePtr New(const String& target_name,
                                  const Array& arguments_descriptor);
 
-  void Insert(const Smi& class_id, const Object& target) const;
+  void EnsureContains(const Smi& class_id, const Object& target) const;
+  ObjectPtr Lookup(const Smi& class_id) const;
 
   void SwitchToBareInstructions();
 
@@ -6858,8 +6859,6 @@
     return RoundedAllocationSize(sizeof(MegamorphicCacheLayout));
   }
 
-  static MegamorphicCachePtr Clone(const MegamorphicCache& from);
-
  private:
   friend class Class;
   friend class MegamorphicCacheTable;
@@ -6867,9 +6866,12 @@
 
   static MegamorphicCachePtr New();
 
-  // The caller must hold Isolate::megamorphic_mutex().
-  void EnsureCapacityLocked() const;
+  // The caller must hold IsolateGroup::type_feedback_mutex().
   void InsertLocked(const Smi& class_id, const Object& target) const;
+  void EnsureCapacityLocked() const;
+  ObjectPtr LookupLocked(const Smi& class_id) const;
+
+  void InsertEntryLocked(const Smi& class_id, const Object& target) const;
 
   static inline void SetEntry(const Array& array,
                               intptr_t index,
diff --git a/runtime/vm/object_test.cc b/runtime/vm/object_test.cc
index 94e3747..ccd43d8 100644
--- a/runtime/vm/object_test.cc
+++ b/runtime/vm/object_test.cc
@@ -3226,6 +3226,57 @@
   EXPECT_EQ(Bool::True().raw(), test_result.raw());
 }
 
+ISOLATE_UNIT_TEST_CASE(MegamorphicCache) {
+  const auto& name = String::Handle(String::New("name"));
+  const auto& args_descriptor =
+      Array::Handle(ArgumentsDescriptor::NewBoxed(1, 1, Object::null_array()));
+
+  const auto& cidA = Smi::Handle(Smi::New(1));
+  const auto& cidB = Smi::Handle(Smi::New(2));
+
+  const auto& valueA = Smi::Handle(Smi::New(42));
+  const auto& valueB = Smi::Handle(Smi::New(43));
+
+  // Test normal insert/lookup methods.
+  {
+    const auto& cache =
+        MegamorphicCache::Handle(MegamorphicCache::New(name, args_descriptor));
+
+    EXPECT(cache.Lookup(cidA) == Object::null());
+    cache.EnsureContains(cidA, valueA);
+    EXPECT(cache.Lookup(cidA) == valueA.raw());
+
+    EXPECT(cache.Lookup(cidB) == Object::null());
+    cache.EnsureContains(cidB, valueB);
+    EXPECT(cache.Lookup(cidB) == valueB.raw());
+  }
+
+  // Try to insert many keys to hit collisions & growth.
+  {
+    const auto& cache =
+        MegamorphicCache::Handle(MegamorphicCache::New(name, args_descriptor));
+
+    auto& cid = Smi::Handle();
+    auto& value = Object::Handle();
+    for (intptr_t i = 0; i < 100; ++i) {
+      cid = Smi::New(100 * i);
+      if (cid.Value() == kIllegalCid) continue;
+
+      value = Smi::New(i);
+      cache.EnsureContains(cid, value);
+    }
+    auto& expected = Object::Handle();
+    for (intptr_t i = 0; i < 100; ++i) {
+      cid = Smi::New(100 * i);
+      if (cid.Value() == kIllegalCid) continue;
+
+      expected = Smi::New(i);
+      value = cache.Lookup(cid);
+      EXPECT(Smi::Cast(value).Equals(Smi::Cast(expected)));
+    }
+  }
+}
+
 ISOLATE_UNIT_TEST_CASE(FieldTests) {
   const String& f = String::Handle(String::New("oneField"));
   const String& getter_f = String::Handle(Field::GetterName(f));
diff --git a/runtime/vm/runtime_entry.cc b/runtime/vm/runtime_entry.cc
index 134fa30..b5f1d00 100644
--- a/runtime/vm/runtime_entry.cc
+++ b/runtime/vm/runtime_entry.cc
@@ -1941,7 +1941,7 @@
 
   // Insert function found into cache.
   const Smi& class_id = Smi::Handle(zone_, Smi::New(cls.id()));
-  data.Insert(class_id, target_function);
+  data.EnsureContains(class_id, target_function);
   arguments_.SetArgAt(0, StubCode::MegamorphicCall());
   arguments_.SetReturn(data);
 }