From e7b2abefabb6a9671e1d30d7af08cd1f32c9a670 Mon Sep 17 00:00:00 2001
From: Neil Dhar <neildhar@meta.com>
Date: Sat, 4 Nov 2023 14:35:02 -0700
Subject: [PATCH] Mitigate WeakMap cycle leak

Summary:
Currently, every `set` operation on a `WeakMap` results in a write to
`self->valueStorage_`, even if the storage was not resized. This means
that a write barrier is executed, and the `valueStorage_` is marked as
a regular cell, instead of going through the special `WeakMap` marking
code in the GC.

This in turn means that if any value in the storage retains a `WeakMap`
key (i.e. a cycle), that cycle will not be collected during a given GC
run. As long as the `WeakMap` is assigned to at least once during each
GC cycle, no such circular references can be collected.

To mitigate this, only write the storage when it is actually resized.
This avoids the write barrier in the common case and allows most
collections to still collect cycles in the `WeakMap`.

This is only a partial fix, since the storage can still be marked
(either through the barrier or when it is in a handle), but it is still
good enough to mitigate most such problems (see #1147).

Reviewed By: tmikov

Differential Revision: D51000231

fbshipit-source-id: 735bac55f19fd9d0a1b58a6d7d77df74c690ab92
---
 lib/VM/JSWeakMapImpl.cpp     |  3 +--
 test/hermes/weakmap-cycle.js | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)
 create mode 100644 test/hermes/weakmap-cycle.js

diff --git a/lib/VM/JSWeakMapImpl.cpp b/lib/VM/JSWeakMapImpl.cpp
index 1760b868969..c1af3ff5b97 100644
--- a/lib/VM/JSWeakMapImpl.cpp
+++ b/lib/VM/JSWeakMapImpl.cpp
@@ -291,6 +291,7 @@ CallResult<uint32_t> JSWeakMapImplBase::getFreeValueStorageIndex(
             ExecutionStatus::EXCEPTION)) {
       return ExecutionStatus::EXCEPTION;
     }
+    self->valueStorage_.setNonNull(runtime, *storageHandle, runtime.getHeap());
   }
 
   // Update internal state here to ensure we don't corrupt it on exception.
@@ -303,8 +304,6 @@ CallResult<uint32_t> JSWeakMapImplBase::getFreeValueStorageIndex(
   }
 
   assert(i < storageHandle->size(runtime) && "invalid index");
-  self->valueStorage_.setNonNull(runtime, *storageHandle, runtime.getHeap());
-
   return i;
 }
 
diff --git a/test/hermes/weakmap-cycle.js b/test/hermes/weakmap-cycle.js
new file mode 100644
index 00000000000..ffea5b8c506
--- /dev/null
+++ b/test/hermes/weakmap-cycle.js
@@ -0,0 +1,18 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// RUN: %hermes -gc-max-heap=32M %s
+// REQUIRES: !slow_debug
+
+(function () {
+    var foo = new WeakMap();
+
+    for (var i = 0; i < 1000000; i++) {
+        var x = {};
+        foo.set(x, x);
+    }
+})();