Reland "[wasm][gc] Free WasmCode objects"

This is a reland of b6fb27077d.

Unchanged reland, TSan issue were fixed in https://crrev.com/c/1593340
and https://crrev.com/c/1594553.

Original change's description:
> [wasm][gc] Free WasmCode objects
>
> This adds the next step to freeing code: We free the actual C++
> {WasmCode} objects. This will cause UAF if any C++ code uses stale
> references.
> The underlying machine code will still not be freed.
>
> For simplicity, this CL changes the vector of owned_code to an ordered
> set, such that lookup and removal is much simpler. The drawback is that
> insertion is now more expensive.
>
> R=mstarzinger@chromium.org
>
> Bug: v8:8217
> Change-Id: I07fc81167816637fbaad6c06ff79e3f952f2fde8
> Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/1593080
> Commit-Queue: Clemens Hammacher <clemensh@chromium.org>
> Reviewed-by: Michael Starzinger <mstarzinger@chromium.org>
> Cr-Commit-Position: refs/heads/master@{#61165}

TBR=mstarzinger@chromium.org

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

src/vector.h

@@ -180,6 +180,12 @@ class OwnedVector {
   constexpr T* begin() const { return start(); }
   constexpr T* end() const { return start() + size(); }
 
+  // Access individual vector elements - checks bounds in debug mode.
+  T& operator[](size_t index) const {
+    DCHECK_LT(index, length_);
+    return data_[index];
+  }
+
   // Returns a {Vector<T>} view of the data in this vector.
   Vector<T> as_vector() const { return Vector<T>(start(), size()); }
 

src/wasm/wasm-code-manager.cc

@@ -420,7 +420,6 @@ NativeModule::NativeModule(WasmEngine* engine, const WasmFeatures& enabled,
       CompilationState::New(*shared_this, std::move(async_counters));
   DCHECK_NOT_NULL(module_);
   owned_code_space_.emplace_back(std::move(code_space));
-  owned_code_.reserve(num_functions());
 
 #if defined(V8_OS_WIN_X64)
   // On some platforms, specifically Win64, we need to reserve some pages at
@@ -786,7 +785,7 @@ WasmCode* NativeModule::PublishCodeLocked(std::unique_ptr<WasmCode> code) {
   }
   WasmCodeRefScope::AddRef(code.get());
   WasmCode* result = code.get();
-  owned_code_.emplace_back(std::move(code));
+  owned_code_.emplace(result->instruction_start(), std::move(code));
   return result;
 }
 
@@ -980,42 +979,14 @@ void NativeModule::SetWireBytes(OwnedVector<const uint8_t> wire_bytes) {
 
 WasmCode* NativeModule::Lookup(Address pc) const {
   base::MutexGuard lock(&allocation_mutex_);
-  if (owned_code_.empty()) return nullptr;
-  // First update the sorted portion counter.
-  if (owned_code_sorted_portion_ == 0) ++owned_code_sorted_portion_;
-  while (owned_code_sorted_portion_ < owned_code_.size() &&
-         owned_code_[owned_code_sorted_portion_ - 1]->instruction_start() <=
-             owned_code_[owned_code_sorted_portion_]->instruction_start()) {
-    ++owned_code_sorted_portion_;
-  }
-  // Execute at most two rounds: First check whether the {pc} is within the
-  // sorted portion of {owned_code_}. If it's not, then sort the whole vector
-  // and retry.
-  while (true) {
-    auto iter =
-        std::upper_bound(owned_code_.begin(), owned_code_.end(), pc,
-                         [](Address pc, const std::unique_ptr<WasmCode>& code) {
-                           DCHECK_NE(kNullAddress, pc);
-                           DCHECK_NOT_NULL(code);
-                           return pc < code->instruction_start();
-                         });
-    if (iter != owned_code_.begin()) {
-      --iter;
-      WasmCode* candidate = iter->get();
-      DCHECK_NOT_NULL(candidate);
-      if (candidate->contains(pc)) {
-        WasmCodeRefScope::AddRef(candidate);
-        return candidate;
-      }
-    }
-    if (owned_code_sorted_portion_ == owned_code_.size()) return nullptr;
-    std::sort(owned_code_.begin(), owned_code_.end(),
-              [](const std::unique_ptr<WasmCode>& code1,
-                 const std::unique_ptr<WasmCode>& code2) {
-                return code1->instruction_start() < code2->instruction_start();
-              });
-    owned_code_sorted_portion_ = owned_code_.size();
-  }
+  auto iter = owned_code_.upper_bound(pc);
+  if (iter == owned_code_.begin()) return nullptr;
+  --iter;
+  WasmCode* candidate = iter->second.get();
+  DCHECK_EQ(candidate->instruction_start(), iter->first);
+  if (!candidate->contains(pc)) return nullptr;
+  WasmCodeRefScope::AddRef(candidate);
+  return candidate;
 }
 
 Address NativeModule::GetCallTargetForFunction(uint32_t func_index) const {
@@ -1376,9 +1347,11 @@ bool NativeModule::IsRedirectedToInterpreter(uint32_t func_index) {
 }
 
 void NativeModule::FreeCode(Vector<WasmCode* const> codes) {
-  // For now, we neither free the {WasmCode} objects, nor do we free any code.
-  // We just zap the code to ensure it's not executed any more.
-  // TODO(clemensh): Actually free the {WasmCode} objects and the code pages.
+  // For now, we only free the {WasmCode} objects and zap the code they referred
+  // to. We do not actually free the code pages yet.
+  // TODO(clemensh): Actually free the underlying code pages.
+
+  // Zap code area.
   size_t code_size = 0;
   for (WasmCode* code : codes) {
     ZapCode(code->instruction_start(), code->instructions().size());
@@ -1387,6 +1360,13 @@ void NativeModule::FreeCode(Vector<WasmCode* const> codes) {
     code_size += code->instructions().size();
   }
   freed_code_size_.fetch_add(code_size);
+
+  // Free the {WasmCode} objects. This will also unregister trap handler data.
+  base::MutexGuard guard(&allocation_mutex_);
+  for (WasmCode* code : codes) {
+    DCHECK_EQ(1, owned_code_.count(code->instruction_start()));
+    owned_code_.erase(code->instruction_start());
+  }
 }
 
 void WasmCodeManager::FreeNativeModule(NativeModule* native_module) {

src/wasm/wasm-code-manager.h

@@ -513,14 +513,9 @@ class V8_EXPORT_PRIVATE NativeModule final {
   //////////////////////////////////////////////////////////////////////////////
   // Protected by {allocation_mutex_}:
 
-  // Holds all allocated code objects. Mutable because it might get sorted in
-  // {Lookup()}.
-  mutable std::vector<std::unique_ptr<WasmCode>> owned_code_;
-
-  // Keep track of the portion of {owned_code_} that is sorted.
-  // Entries [0, owned_code_sorted_portion_) are known to be sorted.
-  // Mutable because it might get modified in {Lookup()}.
-  mutable size_t owned_code_sorted_portion_ = 0;
+  // Holds all allocated code objects. For lookup based on pc, the key is the
+  // instruction start address of the value.
+  std::map<Address, std::unique_ptr<WasmCode>> owned_code_;
 
   std::unique_ptr<WasmCode* []> code_table_;