Reland "[weakrefs] Clear unregister token-related fields when clearing weak cells"

This is a reland of 360c7afca5

Changes since revert:
  - Read the unregister token using a relaxed read during marking

Original change's description:
> [weakrefs] Clear unregister token-related fields when clearing weak cells
>
> Bug: chromium:1213770
> Change-Id: Ic063e79bfa8f3dabdd29d1cc9ed74c7af44d0c31
> Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2923294
> Reviewed-by: Michael Lippautz <mlippautz@chromium.org>
> Commit-Queue: Shu-yu Guo <syg@chromium.org>
> Cr-Commit-Position: refs/heads/master@{#74890}

Bug: chromium:1213770
Change-Id: I8d0b946359b85a4760113e26dbaeaa9479e3b5fd
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2930554
Reviewed-by: Shu-yu Guo <syg@chromium.org>
Reviewed-by: Michael Lippautz <mlippautz@chromium.org>
Commit-Queue: Shu-yu Guo <syg@chromium.org>
Cr-Commit-Position: refs/heads/master@{#74925}

src/builtins/finalization-registry.tq

@@ -143,21 +143,22 @@ FinalizationRegistryRegister(
     ThrowTypeError(
         MessageTemplate::kWeakRefsRegisterTargetAndHoldingsMustNotBeSame);
   }
-  const unregisterToken = arguments[2];
   // 5. If Type(unregisterToken) is not Object,
   //   a. If unregisterToken is not undefined, throw a TypeError exception.
   //   b. Set unregisterToken to empty.
-  let hasUnregisterToken: bool = false;
-  typeswitch (unregisterToken) {
+  const unregisterTokenRaw = arguments[2];
+  let unregisterToken: JSReceiver|Undefined;
+  typeswitch (unregisterTokenRaw) {
     case (Undefined): {
+      unregisterToken = Undefined;
     }
-    case (JSReceiver): {
-      hasUnregisterToken = true;
+    case (unregisterTokenObj: JSReceiver): {
+      unregisterToken = unregisterTokenObj;
     }
     case (JSAny): deferred {
       ThrowTypeError(
           MessageTemplate::kWeakRefsUnregisterTokenMustBeObject,
-          unregisterToken);
+          unregisterTokenRaw);
     }
   }
   // 6. Let cell be the Record { [[WeakRefTarget]] : target, [[HeldValue]]:
@@ -178,7 +179,7 @@ FinalizationRegistryRegister(
   };
   // 7. Append cell to finalizationRegistry.[[Cells]].
   PushCell(finalizationRegistry, cell);
-  if (hasUnregisterToken) {
+  if (unregisterToken != Undefined) {
     // If an unregister token is provided, a runtime call is needed to
     // do some OrderedHashTable operations and register the mapping.
     // See v8:10705.

src/heap/mark-compact.cc

@@ -2561,8 +2561,7 @@ void MarkCompactCollector::ClearJSWeakRefs() {
       RecordSlot(weak_cell, slot, HeapObject::cast(*slot));
     }
 
-    HeapObject unregister_token =
-        HeapObject::cast(weak_cell.unregister_token());
+    HeapObject unregister_token = weak_cell.unregister_token();
     if (!non_atomic_marking_state()->IsBlackOrGrey(unregister_token)) {
       // The unregister token is dead. Remove any corresponding entries in the
       // key map. Multiple WeakCell with the same token will have all their
@@ -2578,11 +2577,6 @@ void MarkCompactCollector::ClearJSWeakRefs() {
             matched_cell.set_unregister_token(undefined);
           },
           gc_notify_updated_slot);
-      // The following is necessary because in the case that weak_cell has
-      // already been popped and removed from the FinalizationRegistry, the call
-      // to JSFinalizationRegistry::RemoveUnregisterToken above will not find
-      // weak_cell itself to clear its unregister token.
-      weak_cell.set_unregister_token(undefined);
     } else {
       // The unregister_token is alive.
       ObjectSlot slot = weak_cell.RawField(WeakCell::kUnregisterTokenOffset);

src/heap/marking-visitor-inl.h

@@ -330,7 +330,7 @@ int MarkingVisitorBase<ConcreteVisitor, MarkingState>::VisitWeakCell(
   this->VisitMapPointer(weak_cell);
   WeakCell::BodyDescriptor::IterateBody(map, weak_cell, size, this);
   HeapObject target = weak_cell.relaxed_target();
-  HeapObject unregister_token = HeapObject::cast(weak_cell.unregister_token());
+  HeapObject unregister_token = weak_cell.relaxed_unregister_token();
   concrete_visitor()->SynchronizePageAccess(target);
   concrete_visitor()->SynchronizePageAccess(unregister_token);
   if (concrete_visitor()->marking_state()->IsBlackOrGrey(target) &&

src/objects/js-weak-refs-inl.h

@@ -163,6 +163,10 @@ HeapObject WeakCell::relaxed_target() const {
   return TaggedField<HeapObject>::Relaxed_Load(*this, kTargetOffset);
 }
 
+HeapObject WeakCell::relaxed_unregister_token() const {
+  return TaggedField<HeapObject>::Relaxed_Load(*this, kUnregisterTokenOffset);
+}
+
 template <typename GCNotifyUpdatedSlotCallback>
 void WeakCell::Nullify(Isolate* isolate,
                        GCNotifyUpdatedSlotCallback gc_notify_updated_slot) {

src/objects/js-weak-refs.h

@@ -93,6 +93,9 @@ class WeakCell : public TorqueGeneratedWeakCell<WeakCell, HeapObject> {
   // Provide relaxed load access to target field.
   inline HeapObject relaxed_target() const;
 
+  // Provide relaxed load access to the unregister token field.
+  inline HeapObject relaxed_unregister_token() const;
+
   // Nullify is called during GC and it modifies the pointers in WeakCell and
   // JSFinalizationRegistry. Thus we need to tell the GC about the modified
   // slots via the gc_notify_updated_slot function. The normal write barrier is

src/objects/js-weak-refs.tq

@@ -22,7 +22,7 @@ extern class JSFinalizationRegistry extends JSObject {
 extern class WeakCell extends HeapObject {
   finalization_registry: Undefined|JSFinalizationRegistry;
   target: Undefined|JSReceiver;
-  unregister_token: JSAny;
+  unregister_token: Undefined|JSReceiver;
   holdings: JSAny;
 
   // For storing doubly linked lists of WeakCells in JSFinalizationRegistry's

src/objects/objects.cc

@@ -6844,6 +6844,7 @@ void JSFinalizationRegistry::RemoveCellFromUnregisterTokenMap(
       JSFinalizationRegistry::cast(Object(raw_finalization_registry));
   WeakCell weak_cell = WeakCell::cast(Object(raw_weak_cell));
   DCHECK(!weak_cell.unregister_token().IsUndefined(isolate));
+  HeapObject undefined = ReadOnlyRoots(isolate).undefined_value();
 
   // Remove weak_cell from the linked list of other WeakCells with the same
   // unregister token and remove its unregister token from key_map if necessary
@@ -6852,7 +6853,7 @@ void JSFinalizationRegistry::RemoveCellFromUnregisterTokenMap(
   if (weak_cell.key_list_prev().IsUndefined(isolate)) {
     SimpleNumberDictionary key_map =
         SimpleNumberDictionary::cast(finalization_registry.key_map());
-    Object unregister_token = weak_cell.unregister_token();
+    HeapObject unregister_token = weak_cell.unregister_token();
     uint32_t key = Smi::ToInt(unregister_token.GetHash());
     InternalIndex entry = key_map.FindEntry(isolate, key);
     DCHECK(entry.is_found());
@@ -6867,8 +6868,7 @@ void JSFinalizationRegistry::RemoveCellFromUnregisterTokenMap(
       // of the key in the hash table.
       WeakCell next = WeakCell::cast(weak_cell.key_list_next());
       DCHECK_EQ(next.key_list_prev(), weak_cell);
-      next.set_key_list_prev(ReadOnlyRoots(isolate).undefined_value());
-      weak_cell.set_key_list_next(ReadOnlyRoots(isolate).undefined_value());
+      next.set_key_list_prev(undefined);
       key_map.ValueAtPut(entry, next);
     }
   } else {
@@ -6880,6 +6880,12 @@ void JSFinalizationRegistry::RemoveCellFromUnregisterTokenMap(
       next.set_key_list_prev(weak_cell.key_list_prev());
     }
   }
+
+  // weak_cell is now removed from the unregister token map, so clear its
+  // unregister token-related fields for heap verification.
+  weak_cell.set_unregister_token(undefined);
+  weak_cell.set_key_list_prev(undefined);
+  weak_cell.set_key_list_next(undefined);
 }
 
 }  // namespace internal

test/cctest/test-js-weak-refs.cc

@@ -815,7 +815,7 @@ TEST(TestRemoveUnregisterToken) {
 
   Handle<JSObject> token1 = CreateKey("token1", isolate);
   Handle<JSObject> token2 = CreateKey("token2", isolate);
-  Handle<Object> undefined =
+  Handle<HeapObject> undefined =
       handle(ReadOnlyRoots(isolate).undefined_value(), isolate);
 
   Handle<WeakCell> weak_cell1a = FinalizationRegistryRegister(
@@ -979,15 +979,17 @@ TEST(UnregisterTokenHeapVerifier) {
   v8::HandleScope outer_scope(isolate);
 
   {
-    // Make a new FinalizationRegistry and register an object with an unregister
-    // token that's unreachable after the IIFE returns.
+    // Make a new FinalizationRegistry and register two objects with the same
+    // unregister token that's unreachable after the IIFE returns.
     v8::HandleScope scope(isolate);
     CompileRun(
         "var token = {}; "
         "var registry = new FinalizationRegistry(function ()  {}); "
         "(function () { "
-        "  let o = {}; "
-        "  registry.register(o, {}, token); "
+        "  let o1 = {}; "
+        "  let o2 = {}; "
+        "  registry.register(o1, {}, token); "
+        "  registry.register(o2, {}, token); "
         "})();");
   }