[wasm][gc] Add data structures for ref-counted WasmCode

This adds support to ref-count uses of WasmCode, and introduces a
{WasmCodeRefScope} to be used whereever WasmCode objects need to be
kept alive, e.g. because a pointer is passed around.
Future CLs will introduce proper scopes in the whole code base and
enable the DCHECK that's currently commented out.

R=titzer@chromium.org

Bug: v8:8217
Change-Id: I1659a0e9d57cd22fe70e6f2661d0d8af9f0906c7
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/1526005
Commit-Queue: Clemens Hammacher <clemensh@chromium.org>
Reviewed-by: Ben Titzer <titzer@chromium.org>
Cr-Commit-Position: refs/heads/master@{#60438}

src/wasm/wasm-code-manager.cc

@@ -1271,6 +1271,10 @@ std::vector<WasmCode*> NativeModule::AddCompiledCode(
   return returned_code;
 }
 
+void NativeModule::FreeCode(Vector<WasmCode* const> codes) {
+  // TODO(clemensh): Implement.
+}
+
 void WasmCodeManager::FreeNativeModule(NativeModule* native_module) {
   base::MutexGuard lock(&native_modules_mutex_);
   TRACE_HEAP("Freeing NativeModule %p\n", native_module);
@@ -1336,6 +1340,66 @@ NativeModuleModificationScope::~NativeModuleModificationScope() {
   }
 }
 
+namespace {
+thread_local WasmCodeRefScope* current_code_refs_scope = nullptr;
+
+// Receives a vector by value which is modified in this function.
+void DecrementRefCount(std::vector<WasmCode*> code_vec) {
+  // Decrement the ref counter of all given code objects. Keep the ones whose
+  // ref count drops to zero.
+  auto remaining_elements_it = code_vec.begin();
+  for (auto it = code_vec.begin(), end = code_vec.end(); it != end; ++it) {
+    if ((*it)->DecRef()) *remaining_elements_it++ = *it;
+  }
+  code_vec.resize(remaining_elements_it - code_vec.begin());
+
+  // Sort the vector by NativeModule, then by instruction start.
+  std::sort(code_vec.begin(), code_vec.end(),
+            [](const WasmCode* a, const WasmCode* b) {
+              return a->native_module() == b->native_module()
+                         ? a->instruction_start() < b->instruction_start()
+                         : a->native_module() < b->native_module();
+            });
+  // For each native module, free all its code objects at once.
+  auto range_begin = code_vec.begin();
+  while (range_begin != code_vec.end()) {
+    NativeModule* native_module = (*range_begin)->native_module();
+    auto range_end = range_begin + 1;
+    while (range_end < code_vec.end() &&
+           (*range_end)->native_module() == native_module) {
+      ++range_end;
+    }
+    size_t range_size = static_cast<size_t>(range_end - range_begin);
+    Vector<WasmCode*> code_vec{&*range_begin, range_size};
+    native_module->FreeCode(code_vec);
+    range_begin = range_end;
+  }
+}
+
+}  // namespace
+
+WasmCodeRefScope::WasmCodeRefScope()
+    : previous_scope_(current_code_refs_scope) {
+  current_code_refs_scope = this;
+}
+
+WasmCodeRefScope::~WasmCodeRefScope() {
+  DCHECK_EQ(this, current_code_refs_scope);
+  current_code_refs_scope = previous_scope_;
+  DecrementRefCount({code_ptrs_.begin(), code_ptrs_.end()});
+}
+
+// static
+void WasmCodeRefScope::AddRef(WasmCode* code) {
+  WasmCodeRefScope* current_scope = current_code_refs_scope;
+  // TODO(clemensh): Remove early return, activate DCHECK instead.
+  // DCHECK_NOT_NULL(current_scope);
+  if (!current_scope) return;
+  auto entry = current_scope->code_ptrs_.insert(code);
+  // If we added a new entry, increment the ref counter.
+  if (entry.second) code->IncRef();
+}
+
 }  // namespace wasm
 }  // namespace internal
 }  // namespace v8

src/wasm/wasm-code-manager.h

@@ -9,6 +9,7 @@
 #include <list>
 #include <map>
 #include <memory>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -146,6 +147,19 @@ class V8_EXPORT_PRIVATE WasmCode final {
 
   ~WasmCode();
 
+  void IncRef() {
+    int old_val = ref_count_.fetch_add(1, std::memory_order_relaxed);
+    DCHECK_LE(1, old_val);
+    DCHECK_GT(kMaxInt, old_val);
+    USE(old_val);
+  }
+
+  V8_WARN_UNUSED_RESULT bool DecRef() {
+    int old_count = ref_count_.fetch_sub(1, std::memory_order_relaxed);
+    DCHECK_LE(1, old_count);
+    return old_count == 1;
+  }
+
   enum FlushICache : bool { kFlushICache = true, kNoFlushICache = false };
 
   static constexpr uint32_t kAnonymousFuncIndex = 0xffffffff;
@@ -216,6 +230,18 @@ class V8_EXPORT_PRIVATE WasmCode final {
   OwnedVector<trap_handler::ProtectedInstructionData> protected_instructions_;
   Tier tier_;
 
+  // WasmCode is ref counted. Counters are held by:
+  //   1) The jump table.
+  //   2) Function tables.
+  //   3) {WasmCodeRefScope}s.
+  //   4) Threads currently executing this code.
+  // If a decrement of (1) or (2) would drop the ref count to 0, that code
+  // becomes a candidate for garbage collection. At that point, we add
+  // ref counts for (4) *before* decrementing the counter to ensure the code
+  // stays alive as long as it's being used. Once the ref count drops to zero,
+  // the code object is deleted and the memory for the machine code is freed.
+  std::atomic<int> ref_count_{1};
+
   DISALLOW_COPY_AND_ASSIGN(WasmCode);
 };
 
@@ -371,6 +397,11 @@ class V8_EXPORT_PRIVATE NativeModule final {
   WasmCode* AddCompiledCode(WasmCompilationResult);
   std::vector<WasmCode*> AddCompiledCode(Vector<WasmCompilationResult>);
 
+  // Free a set of functions of this module. Uncommits whole pages if possible.
+  // The given vector must be ordered by the instruction start address, and all
+  // {WasmCode} objects must not be used any more.
+  void FreeCode(Vector<WasmCode* const>);
+
  private:
   friend class WasmCode;
   friend class WasmCodeManager;
@@ -570,6 +601,25 @@ class NativeModuleModificationScope final {
   NativeModule* native_module_;
 };
 
+// {WasmCodeRefScope}s form a perfect stack. New {WasmCode} pointers generated
+// by e.g. creating new code or looking up code by its address are added to the
+// top-most {WasmCodeRefScope}.
+class V8_EXPORT_PRIVATE WasmCodeRefScope {
+ public:
+  WasmCodeRefScope();
+  ~WasmCodeRefScope();
+
+  // Register a {WasmCode} reference in the current {WasmCodeRefScope}. Fails if
+  // there is no current scope.
+  static void AddRef(WasmCode*);
+
+ private:
+  WasmCodeRefScope* const previous_scope_;
+  std::unordered_set<WasmCode*> code_ptrs_;
+
+  DISALLOW_COPY_AND_ASSIGN(WasmCodeRefScope);
+};
+
 }  // namespace wasm
 }  // namespace internal
 }  // namespace v8