Don't reset all inline cache types when one changes

Summary:
Clearing the whole cache seems counter productive... we seem to be doing this because we also watch the descriptor type when it's a data descriptor. Instead let's check to see if we have a data descriptor and invalidate that type when the descriptor changes.

This is a little more complicated than it seems because of finalization at shutdown. The descriptor could be freed as part of cyclic trash. Worse still the descriptor could live on a base type of the type that we're freeing. And worse still is that the weak references tracked in `tp_subclasses` could be cleared on the base type before the base type is cleared causing us not to receive the type modification event for the sub type. For this reason we need to track the descriptors type and not rely on doing `Py_TYPE(descr)` when we're finalizing our caches.

Reviewed By: alexmalyshev

Differential Revision: D104697182

fbshipit-source-id: 25d421b436bc91538df54d245edf50b71a27b648

Author: DinoV

cinderx/Jit/inline_cache.cpp

@@ -27,6 +27,10 @@ struct TypeWatcher {
   jit::UnorderedMap<BorrowedRef<PyTypeObject>, jit::UnorderedSet<T*>> caches;
 
   void watch(BorrowedRef<PyTypeObject> type, T* cache) {
+    if (PyType_HasFeature(type, Py_TPFLAGS_IMMUTABLETYPE)) {
+      // We don't watch immutable types - they can't be modified.
+      return;
+    }
     JIT_CHECK(
         cinderx::getModuleState()->watcher_state.watchType(type) == 0,
         "Failed to watch type {} for attribute cache",
@@ -43,20 +47,30 @@ struct TypeWatcher {
     // don't unwatch type; other watchers may still be watching it
   }
 
-  void typeChanged(BorrowedRef<PyTypeObject> type) {
+  // Notify caches watching a type. Calls Callback(cache, type) for each
+  // registered cache. The default callback calls cache->typeChanged(type).
+  template <typename Callback>
+  void typeChanged(BorrowedRef<PyTypeObject> type, Callback cb) {
     auto it = caches.find(type);
     if (it == caches.end()) {
       return;
     }
     jit::UnorderedSet<T*> to_notify = std::move(it->second);
     caches.erase(it);
     for (T* cache : to_notify) {
-      cache->typeChanged(type);
+      cb(cache, type);
     }
   }
+
+  void typeChanged(BorrowedRef<PyTypeObject> type) {
+    typeChanged(type, [](T* cache, BorrowedRef<PyTypeObject> tp) {
+      cache->typeChanged(tp);
+    });
+  }
 };
 
 TypeWatcher<AttributeCache> ac_watcher;
+TypeWatcher<AttributeCache> ac_descr_watcher;
 TypeWatcher<LoadTypeAttrCache> ltac_watcher;
 TypeWatcher<LoadMethodCache> lm_watcher;
 TypeWatcher<LoadTypeMethodCache> ltm_watcher;
@@ -563,6 +577,7 @@ void AttributeMutator::set_combined(PyTypeObject* type) {
 void AttributeMutator::set_data_descr(PyTypeObject* type, PyObject* descr) {
   set_type(type, Kind::kDataDescr);
   data_descr_.descr = descr;
+  data_descr_.descr_type = Py_TYPE(descr);
 }
 
 void AttributeMutator::set_member_descr(PyTypeObject* type, PyObject* descr) {
@@ -589,6 +604,13 @@ void AttributeMutator::set_split(
   split_.keys = keys;
 }
 
+BorrowedRef<PyTypeObject> AttributeMutator::watchedDescrType() const {
+  if (get_kind() == Kind::kDataDescr) {
+    return data_descr_.descr_type;
+  }
+  return nullptr;
+}
+
 inline int
 AttributeMutator::setAttr(PyObject* obj, PyObject* name, PyObject* value) {
   JIT_CHECK(
@@ -675,14 +697,61 @@ AttributeCache::~AttributeCache() {
   for (auto& entry : entries()) {
     if (entry.type() != nullptr) {
       ac_watcher.unwatch(entry.type(), this);
+      BorrowedRef<PyTypeObject> descr_tp = entry.watchedDescrType();
+      if (descr_tp != nullptr) {
+        ac_descr_watcher.unwatch(descr_tp, this);
+      }
       entry.reset();
     }
   }
 }
 
-void AttributeCache::typeChanged(PyTypeObject*) {
+void AttributeCache::typeChanged(PyTypeObject* tp) {
   for (auto& entry : entries()) {
-    entry.reset();
+    if (entry.type() == tp) {
+      BorrowedRef<PyTypeObject> descr_tp = entry.watchedDescrType();
+      // Reset the entry, this also resets the kind so after the call
+      // watchedDescrType returns nullptr.
+      entry.reset();
+      if (descr_tp != nullptr) {
+        // For entry.type() the cache is 1 to 1 between types and the cache.
+        // For descriptors the same descriptor can be shared multiple times
+        // per a type. So check if we have any descriptors remaining and only
+        // unwatch when they're all clear.
+        bool found = false;
+        for (auto& other : entries()) {
+          if (other.watchedDescrType() == descr_tp) {
+            found = true;
+            break;
+          }
+        }
+        if (!found) {
+          // there are no other entries watching this descriptor.
+          ac_descr_watcher.unwatch(descr_tp, this);
+        }
+      }
+    }
+  }
+}
+
+void AttributeCache::descrTypeChanged(PyTypeObject* tp) {
+  // We only need to unwatch this AttributeCache in the ac_descr_watcher
+  // once as it applies to all entries so we use a flag to track if we've
+  // already unwatched it for multiple entries (it wouldn't hurt to try
+  // and unwatch multiple times but it's a couple of hash lookups)
+  bool found = false;
+  for (auto& entry : entries()) {
+    if (entry.watchedDescrType() == tp) {
+      // If we were watching the descriptor the whole cache entry is
+      // invalidated and we no longer need to watch the type.
+      ac_watcher.unwatch(entry.type(), this);
+      entry.reset();
+      if (!found) {
+        // unwatch the descriptor watcher once for the whole cache
+        ac_descr_watcher.unwatch(tp, this);
+        found = true;
+      }
+    }
   }
 }
 
@@ -766,9 +835,10 @@ void AttributeCache::fill(
       if (descr_type == &PyMemberDescr_Type) {
         mut->set_member_descr(type, descr);
       } else {
-        // If someone deletes descr_types's __set__ method, it will no longer
-        // be a data descriptor, and the cache kind has to change.
-        ac_watcher.watch(descr_type, this);
+        // If someone modifies descr_type (e.g., deletes __set__), it may no
+        // longer be a data descriptor. Watch it via the descriptor watcher
+        // so the cache is invalidated.
+        ac_descr_watcher.watch(descr_type, this);
         mut->set_data_descr(type, descr);
       }
     } else {
@@ -1602,6 +1672,10 @@ LoadModuleMethodCache::lookupSlowPath(BorrowedRef<> obj, BorrowedRef<> name) {
 
 void notifyICsTypeChanged(BorrowedRef<PyTypeObject> type) {
   ac_watcher.typeChanged(type);
+  ac_descr_watcher.typeChanged(
+      type, [](AttributeCache* cache, BorrowedRef<PyTypeObject> tp) {
+        cache->descrTypeChanged(tp);
+      });
   ltac_watcher.typeChanged(type);
   lm_watcher.typeChanged(type);
   ltm_watcher.typeChanged(type);

cinderx/Jit/inline_cache.h

@@ -55,6 +55,7 @@ struct DataDescrMutator {
   int setAttr(PyObject* obj, PyObject* value);
 
   BorrowedRef<> descr;
+  BorrowedRef<PyTypeObject> descr_type;
 };
 
 // Mutator for a member descriptor
@@ -111,6 +112,7 @@ class AttributeMutator {
       Py_ssize_t val_offset,
       PyDictKeysObject* keys,
       bool values_inline);
+  BorrowedRef<PyTypeObject> watchedDescrType() const;
 
   PyObject* getAttr(PyObject* obj, PyObject* name);
   int setAttr(PyObject* obj, PyObject* name, PyObject* value);
@@ -139,6 +141,7 @@ class AttributeCache {
   ~AttributeCache();
 
   void typeChanged(PyTypeObject* type);
+  void descrTypeChanged(PyTypeObject* type);
 
  protected:
   std::span<AttributeMutator> entries();

cinderx/PythonLib/test_cinderx/test_jit_attr_cache.py

@@ -430,6 +430,59 @@ def get_attr(o):
         C.__dict__["foo"].__class__ = Descr
         self.assertEqual(get_attr(c), "get")
 
+    def test_shared_descr_type_across_cache_entries(self):
+        """Two types share the same descriptor type in the same cache.
+
+        When one type changes and its cache entry is reset, the descriptor
+        type watch must survive for the remaining entry. Otherwise mutating
+        the descriptor type (e.g. removing __set__) won't invalidate the
+        surviving entry.
+        """
+
+        class Descr:
+            def __get__(self, obj, ty):
+                return "descr"
+
+            def __set__(self, obj, val):
+                raise RuntimeError("unimplemented")
+
+        class T1:
+            foo = Descr()
+
+        class T2:
+            foo = Descr()
+
+        @cinder_support.failUnlessJITCompiled
+        def get_attr(o):
+            return o.foo
+
+        t1 = T1()
+        t2 = T2()
+
+        # Prime and cache both entries (same descriptor type Descr).
+        self.assertEqual(get_attr(t1), "descr")
+        self.assertEqual(get_attr(t1), "descr")
+        self.assertEqual(get_attr(t2), "descr")
+        self.assertEqual(get_attr(t2), "descr")
+
+        # Both t1 and t2 have instance dict entries shadowed by
+        # the data descriptor.
+        t1.__dict__["foo"] = "t1 attr"
+        t2.__dict__["foo"] = "t2 attr"
+        self.assertEqual(get_attr(t1), "descr")
+        self.assertEqual(get_attr(t2), "descr")
+
+        # Invalidate T1's cache entry by mutating T1. This must NOT
+        # unwatch Descr from the descriptor watcher since T2's entry
+        # still depends on it.
+        T1.bar = 1
+
+        # Now mutate the descriptor type: removing __set__ makes Descr
+        # no longer a data descriptor. T2's entry must be invalidated
+        # so the instance attribute is returned instead.
+        del Descr.__set__
+        self.assertEqual(get_attr(t2), "t2 attr")
+
     @passIf(
         cinderx.jit.is_enabled(),
         "T214641462: Not clear why this is failing, but it is",