Bug 1758223 - Use atomic ops to read / write node flags from stylo. r=nika

The flags stylo cares about reading and writing potentially at the same
time are disjoint, so there's no need for any strong memory ordering.

Differential Revision: https://phabricator.services.mozilla.com/D141829

dom/base/nsWrapperCache.h

@@ -9,6 +9,8 @@
 
 #include "nsCycleCollectionParticipant.h"
 #include "mozilla/Assertions.h"
+#include "mozilla/ServoUtils.h"
+#include "mozilla/RustCell.h"
 #include "js/HeapAPI.h"
 #include "js/TracingAPI.h"
 #include "js/TypeDecls.h"
@@ -89,15 +91,7 @@ class JS_HAZ_ROOTED nsWrapperCache {
  public:
   NS_DECLARE_STATIC_IID_ACCESSOR(NS_WRAPPERCACHE_IID)
 
-  nsWrapperCache()
-      : mWrapper(nullptr),
-        mFlags(0)
-#ifdef BOOL_FLAGS_ON_WRAPPER_CACHE
-        ,
-        mBoolFlags(0)
-#endif
-  {
-  }
+  nsWrapperCache() = default;
   ~nsWrapperCache() {
     // Preserved wrappers should never end up getting cleared, but this can
     // happen during shutdown when a leaked wrapper object is finalized, causing
@@ -256,33 +250,43 @@ class JS_HAZ_ROOTED nsWrapperCache {
 
   using FlagsType = uint32_t;
 
-  FlagsType GetFlags() const { return mFlags & ~kWrapperFlagsMask; }
+  FlagsType GetFlags() const {
+    MOZ_ASSERT(NS_IsMainThread());
+    MOZ_ASSERT(!mozilla::IsInServoTraversal());
+    return mFlags.Get() & ~kWrapperFlagsMask;
+  }
 
+  // This can be called from stylo threads too, so it needs to be atomic, as
+  // this value may be mutated from multiple threads during servo traversal from
+  // rust.
   bool HasFlag(FlagsType aFlag) const {
     MOZ_ASSERT((aFlag & kWrapperFlagsMask) == 0, "Bad flag mask");
-    return !!(mFlags & aFlag);
+    return __atomic_load_n(mFlags.AsPtr(), __ATOMIC_RELAXED) & aFlag;
   }
 
   // Identical to HasFlag, but more explicit about its handling of multiple
-  // flags.
-  bool HasAnyOfFlags(FlagsType aFlags) const {
-    MOZ_ASSERT((aFlags & kWrapperFlagsMask) == 0, "Bad flag mask");
-    return !!(mFlags & aFlags);
-  }
+  // flags. This can be called from stylo threads too.
+  bool HasAnyOfFlags(FlagsType aFlags) const { return HasFlag(aFlags); }
 
+  // This can also be called from stylo, in the sequential part of the
+  // traversal, though it's probably not worth differentiating them for the
+  // purposes of assertions.
   bool HasAllFlags(FlagsType aFlags) const {
     MOZ_ASSERT((aFlags & kWrapperFlagsMask) == 0, "Bad flag mask");
-    return (mFlags & aFlags) == aFlags;
+    return (__atomic_load_n(mFlags.AsPtr(), __ATOMIC_RELAXED) & aFlags) ==
+           aFlags;
   }
 
   void SetFlags(FlagsType aFlagsToSet) {
+    MOZ_ASSERT(NS_IsMainThread());
     MOZ_ASSERT((aFlagsToSet & kWrapperFlagsMask) == 0, "Bad flag mask");
-    mFlags |= aFlagsToSet;
+    mFlags.Set(mFlags.Get() | aFlagsToSet);
   }
 
   void UnsetFlags(FlagsType aFlagsToUnset) {
+    MOZ_ASSERT(NS_IsMainThread());
     MOZ_ASSERT((aFlagsToUnset & kWrapperFlagsMask) == 0, "Bad flag mask");
-    mFlags &= ~aFlagsToUnset;
+    mFlags.Set(mFlags.Get() & ~aFlagsToUnset);
   }
 
   void PreserveWrapper(nsISupports* aScriptObjectHolder) {
@@ -335,23 +339,31 @@ class JS_HAZ_ROOTED nsWrapperCache {
  private:
   void SetWrapperJSObject(JSObject* aWrapper);
 
-  FlagsType GetWrapperFlags() const { return mFlags & kWrapperFlagsMask; }
+  FlagsType GetWrapperFlags() const {
+    // These are used from workers and from the main thread (but shouldn't be
+    // used on the main thread during the Servo traversal).
+    MOZ_ASSERT(!NS_IsMainThread() || !mozilla::IsInServoTraversal());
+    return mFlags.Get() & kWrapperFlagsMask;
+  }
 
   bool HasWrapperFlag(FlagsType aFlag) const {
+    MOZ_ASSERT(!NS_IsMainThread() || !mozilla::IsInServoTraversal());
     MOZ_ASSERT((aFlag & ~kWrapperFlagsMask) == 0, "Bad wrapper flag bits");
-    return !!(mFlags & aFlag);
+    return !!(mFlags.Get() & aFlag);
   }
 
   void SetWrapperFlags(FlagsType aFlagsToSet) {
+    MOZ_ASSERT(!NS_IsMainThread() || !mozilla::IsInServoTraversal());
     MOZ_ASSERT((aFlagsToSet & ~kWrapperFlagsMask) == 0,
                "Bad wrapper flag bits");
-    mFlags |= aFlagsToSet;
+    mFlags.Set(mFlags.Get() | aFlagsToSet);
   }
 
   void UnsetWrapperFlags(FlagsType aFlagsToUnset) {
+    MOZ_ASSERT(!NS_IsMainThread() || !mozilla::IsInServoTraversal());
     MOZ_ASSERT((aFlagsToUnset & ~kWrapperFlagsMask) == 0,
                "Bad wrapper flag bits");
-    mFlags &= ~aFlagsToUnset;
+    mFlags.Set(mFlags.Get() & ~aFlagsToUnset);
   }
 
   void HoldJSObjects(void* aScriptObjectHolder, nsScriptObjectTracer* aTracer,
@@ -379,12 +391,22 @@ class JS_HAZ_ROOTED nsWrapperCache {
 
   enum { kWrapperFlagsMask = WRAPPER_BIT_PRESERVED };
 
-  JSObject* mWrapper;
-  FlagsType mFlags;
+  JSObject* mWrapper = nullptr;
+
+  // Rust code needs to read and write some flags atomically, but we don't want
+  // to make the wrapper flags atomic whole-sale because main-thread code would
+  // become more expensive (loads wouldn't change, but flag setting /
+  // unsetting could become slower enough to be noticeable). Making this an
+  // Atomic whole-sale needs more measuring.
+  //
+  // In order to not mess with aliasing rules the type should not be frozen, so
+  // we use a RustCell, which contains an UnsafeCell internally. See also the
+  // handling of ServoData (though that's a bit different).
+  mozilla::RustCell<FlagsType> mFlags{0};
 
  protected:
 #ifdef BOOL_FLAGS_ON_WRAPPER_CACHE
-  uint32_t mBoolFlags;
+  uint32_t mBoolFlags = 0;
 #endif
 };
 

js/src/devtools/rootAnalysis/analyzeHeapWrites.js

@@ -146,10 +146,6 @@ function treatAsSafeArgument(entry, varName, csuName)
         [/^Gecko_/, null, "nsStyleImageLayers"],
         [/^Gecko_/, null, /FontFamilyList/],
 
-        // RawGeckoBorrowedNode thread-mutable parameters.
-        ["Gecko_SetNodeFlags", "aNode", null],
-        ["Gecko_UnsetNodeFlags", "aNode", null],
-
         // Various Servo binding out parameters. This is a mess and there needs
         // to be a way to indicate which params are out parameters, either using
         // an attribute or a naming convention.

layout/style/GeckoBindings.cpp

@@ -261,15 +261,6 @@ bool Gecko_IsRootElement(const Element* aElement) {
   return aElement->OwnerDoc()->GetRootElement() == aElement;
 }
 
-// Dirtiness tracking.
-void Gecko_SetNodeFlags(const nsINode* aNode, uint32_t aFlags) {
-  const_cast<nsINode*>(aNode)->SetFlags(aFlags);
-}
-
-void Gecko_UnsetNodeFlags(const nsINode* aNode, uint32_t aFlags) {
-  const_cast<nsINode*>(aNode)->UnsetFlags(aFlags);
-}
-
 void Gecko_NoteDirtyElement(const Element* aElement) {
   MOZ_ASSERT(NS_IsMainThread());
   const_cast<Element*>(aElement)->NoteDirtyForServo();

layout/style/GeckoBindings.h

@@ -321,8 +321,6 @@ void Gecko_SetListStyleImageImageValue(
 void Gecko_CopyListStyleImageFrom(nsStyleList* dest, const nsStyleList* src);
 
 // Dirtiness tracking.
-void Gecko_SetNodeFlags(const nsINode* node, uint32_t flags);
-void Gecko_UnsetNodeFlags(const nsINode* node, uint32_t flags);
 void Gecko_NoteDirtyElement(const mozilla::dom::Element*);
 void Gecko_NoteDirtySubtreeForInvalidation(const mozilla::dom::Element*);
 void Gecko_NoteAnimationOnlyDirtyElement(const mozilla::dom::Element*);

layout/style/RustCell.h

@@ -25,8 +25,16 @@ template <typename T>
 class RustCell {
  public:
   RustCell() : mValue() {}
+  explicit RustCell(T aValue) : mValue(aValue) {}
 
   T Get() const { return mValue; }
+  void Set(T aValue) { mValue = aValue; }
+
+  // That this API doesn't mirror the API of rust's Cell because all values in
+  // C++ effectively act like they're wrapped in Cell<...> already, so this type
+  // only exists for FFI purposes.
+  T* AsPtr() { return &mValue; }
+  const T* AsPtr() const { return &mValue; }
 
  private:
   T mValue;

servo/components/style/gecko/wrapper.rs

@@ -80,6 +80,7 @@ use selectors::matching::{ElementSelectorFlags, MatchingContext};
 use selectors::sink::Push;
 use selectors::{Element, OpaqueElement};
 use servo_arc::{Arc, ArcBorrow, RawOffsetArc};
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::cell::RefCell;
 use std::fmt;
 use std::hash::{Hash, Hasher};
@@ -270,9 +271,29 @@ impl<'ln> GeckoNode<'ln> {
         GeckoNode(&content._base)
     }
 
+    #[inline]
+    fn flags_atomic(&self) -> &AtomicU32 {
+        use std::cell::Cell;
+        let flags: &Cell<u32> = &(self.0)._base._base_1.mFlags;
+
+        #[allow(dead_code)]
+        fn static_assert() {
+            let _: [u8; std::mem::size_of::<Cell<u32>>()] = [0u8; std::mem::size_of::<AtomicU32>()];
+            let _: [u8; std::mem::align_of::<Cell<u32>>()] = [0u8; std::mem::align_of::<AtomicU32>()];
+        }
+
+        // Rust doesn't provide standalone atomic functions like GCC/clang do
+        // (via the atomic intrinsics) or via std::atomic_ref, but it guarantees
+        // that the memory representation of u32 and AtomicU32 matches:
+        // https://doc.rust-lang.org/std/sync/atomic/struct.AtomicU32.html
+        unsafe {
+            std::mem::transmute::<&Cell<u32>, &AtomicU32>(flags)
+        }
+    }
+
     #[inline]
     fn flags(&self) -> u32 {
-        (self.0)._base._base_1.mFlags
+        self.flags_atomic().load(Ordering::Relaxed)
     }
 
     #[inline]
@@ -653,18 +674,14 @@ impl<'le> GeckoElement<'le> {
         self.as_node().flags()
     }
 
-    // FIXME: We can implement this without OOL calls, but we can't easily given
-    // GeckoNode is a raw reference.
-    //
-    // We can use a Cell<T>, but that's a bit of a pain.
     #[inline]
     fn set_flags(&self, flags: u32) {
-        unsafe { Gecko_SetNodeFlags(self.as_node().0, flags) }
+        self.as_node().flags_atomic().fetch_or(flags, Ordering::Relaxed);
     }
 
     #[inline]
     unsafe fn unset_flags(&self, flags: u32) {
-        Gecko_UnsetNodeFlags(self.as_node().0, flags)
+        self.as_node().flags_atomic().fetch_and(!flags, Ordering::Relaxed);
     }
 
     /// Returns true if this element has descendants for lazy frame construction.