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

@@ -12,16 +12,9 @@ namespace atom {
 
 namespace {
 
-struct ObjectKey {
+void OnObjectGC(const v8::WeakCallbackInfo<KeyWeakMap::KeyObject>& data) {
-  ObjectKey(int id, KeyWeakMap* map) : id(id), map(map) {}
+  KeyWeakMap::KeyObject* key_object = data.GetParameter();
-  int id;
+  key_object->self->Remove(key_object->key);
-  KeyWeakMap* map;
-};
-
-void OnObjectGC(const v8::WeakCallbackInfo<ObjectKey>& data) {
-  ObjectKey* key = data.GetParameter();
-  key->map->Remove(key->id);
-  delete 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));
+  auto value = make_linked_ptr(new v8::Global<v8::Object>(isolate, object));
-  ObjectKey* key = new ObjectKey(id, this);
+  KeyObject key_object = {key, this};
-  global->SetWeak(key, OnObjectGC, v8::WeakCallbackType::kParameter);
+  auto& p = map_[key] = std::make_pair(key_object, value);
-  map_[id] = global;
+  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_)
+  for (const auto& iter : map_) {
-    keys.emplace_back(v8::Local<v8::Object>::New(isolate, *iter.second));
+    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
+    return;
-    map_.erase(iter);
+  }
+
+  iter->second.second->ClearWeak();
+  map_.erase(iter);
 }
 
 IDWeakMap::IDWeakMap() : next_id_(0) {

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);
 };