From e6c0b1fe0c13c59e0c783dd4c0d6cf7a88b0b859 Mon Sep 17 00:00:00 2001
From: Cheng Zhao <zcbenz@gmail.com>
Date: Wed, 11 May 2016 14:26:32 +0900
Subject: [PATCH] Fix leak when KeyWeakMap::Remove is called directly

Usually the KeyObject would be destroyed when GC happens, but then
Remove is called before GC happens, the KeyObject would be leaked
forever. This fixes it by keeping KeyObject as a member of map.
---
 atom/common/id_weak_map.cc | 42 +++++++++++++++++++-------------------
 atom/common/id_weak_map.h  | 10 ++++++++-
 2 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/atom/common/id_weak_map.cc b/atom/common/id_weak_map.cc
index 2067585eb8f5..c925f1f055c5 100644
--- a/atom/common/id_weak_map.cc
+++ b/atom/common/id_weak_map.cc
@@ -12,16 +12,9 @@ namespace atom {
 
 namespace {
 
-struct ObjectKey {
-  ObjectKey(int id, KeyWeakMap* map) : id(id), map(map) {}
-  int id;
-  KeyWeakMap* map;
-};
-
-void OnObjectGC(const v8::WeakCallbackInfo<ObjectKey>& data) {
-  ObjectKey* key = data.GetParameter();
-  key->map->Remove(key->id);
-  delete key;
+void OnObjectGC(const v8::WeakCallbackInfo<KeyWeakMap::KeyObject>& data) {
+  KeyWeakMap::KeyObject* key_object = data.GetParameter();
+  key_object->self->Remove(key_object->key);
 }
 
 }  // namespace
@@ -30,15 +23,17 @@ KeyWeakMap::KeyWeakMap() {
 }
 
 KeyWeakMap::~KeyWeakMap() {
+  for (const auto& p : map_)
+    p.second.second->ClearWeak();
 }
 
 void KeyWeakMap::Set(v8::Isolate* isolate,
-                    int32_t id,
+                    int32_t key,
                     v8::Local<v8::Object> object) {
-  auto global = make_linked_ptr(new v8::Global<v8::Object>(isolate, object));
-  ObjectKey* key = new ObjectKey(id, this);
-  global->SetWeak(key, OnObjectGC, v8::WeakCallbackType::kParameter);
-  map_[id] = global;
+  auto value = make_linked_ptr(new v8::Global<v8::Object>(isolate, object));
+  KeyObject key_object = {key, this};
+  auto& p = map_[key] = std::make_pair(key_object, value);
+  value->SetWeak(&(p.first), OnObjectGC, v8::WeakCallbackType::kParameter);
 }
 
 v8::MaybeLocal<v8::Object> KeyWeakMap::Get(v8::Isolate* isolate, int32_t id) {
@@ -46,7 +41,7 @@ v8::MaybeLocal<v8::Object> KeyWeakMap::Get(v8::Isolate* isolate, int32_t id) {
   if (iter == map_.end())
     return v8::MaybeLocal<v8::Object>();
   else
-    return v8::Local<v8::Object>::New(isolate, *iter->second);
+    return v8::Local<v8::Object>::New(isolate, *(iter->second.second));
 }
 
 bool KeyWeakMap::Has(int32_t id) const {
@@ -56,17 +51,22 @@ bool KeyWeakMap::Has(int32_t id) const {
 std::vector<v8::Local<v8::Object>> KeyWeakMap::Values(v8::Isolate* isolate) {
   std::vector<v8::Local<v8::Object>> keys;
   keys.reserve(map_.size());
-  for (const auto& iter : map_)
-    keys.emplace_back(v8::Local<v8::Object>::New(isolate, *iter.second));
+  for (const auto& iter : map_) {
+    const auto& value = *(iter.second.second);
+    keys.emplace_back(v8::Local<v8::Object>::New(isolate, value));
+  }
   return keys;
 }
 
 void KeyWeakMap::Remove(int32_t id) {
   auto iter = map_.find(id);
-  if (iter == map_.end())
+  if (iter == map_.end()) {
     LOG(WARNING) << "Removing unexist object with ID " << id;
-  else
-    map_.erase(iter);
+    return;
+  }
+
+  iter->second.second->ClearWeak();
+  map_.erase(iter);
 }
 
 IDWeakMap::IDWeakMap() : next_id_(0) {
diff --git a/atom/common/id_weak_map.h b/atom/common/id_weak_map.h
index a3c49b2346c7..f2014b1cd7fd 100644
--- a/atom/common/id_weak_map.h
+++ b/atom/common/id_weak_map.h
@@ -16,6 +16,12 @@ namespace atom {
 // Like ES6's WeakMap, but the key is Integer and the value is Weak Pointer.
 class KeyWeakMap {
  public:
+  // Records the key and self, used by SetWeak.
+  struct KeyObject {
+    int32_t key;
+    KeyWeakMap* self;
+  };
+
   KeyWeakMap();
   virtual ~KeyWeakMap();
 
@@ -36,7 +42,9 @@ class KeyWeakMap {
 
  private:
   // Map of stored objects.
-  std::unordered_map<int32_t, linked_ptr<v8::Global<v8::Object>>> map_;
+  std::unordered_map<
+      int32_t,
+      std::pair<KeyObject, linked_ptr<v8::Global<v8::Object>>>> map_;
 
   DISALLOW_COPY_AND_ASSIGN(KeyWeakMap);
 };