Make each class only have one weak map

atom/browser/api/atom_api_session.h

@@ -19,7 +19,7 @@ class AtomBrowserContext;
 
 namespace api {
 
-class Session: public mate::TrackableObject {
+class Session: public mate::TrackableObject<Session> {
  public:
   using ResolveProxyCallback = base::Callback<void(std::string)>;
 

atom/browser/api/atom_api_web_contents.cc

@@ -424,6 +424,7 @@ void WebContents::WebContentsDestroyed() {
   // The RenderViewDeleted was not called when the WebContents is destroyed.
   RenderViewDeleted(web_contents()->GetRenderViewHost());
   Emit("destroyed");
+  RemoveFromWeakMap();
 }
 
 void WebContents::NavigationEntryCommitted(

atom/browser/api/atom_api_web_contents.h

@@ -46,7 +46,7 @@ struct SetSizeParams {
   scoped_ptr<gfx::Size> normal_size;
 };
 
-class WebContents : public mate::TrackableObject,
+class WebContents : public mate::TrackableObject<WebContents>,
                     public content::BrowserPluginGuestDelegate,
                     public CommonWebContentsDelegate,
                     public content::WebContentsObserver,

atom/browser/api/atom_api_window.cc

@@ -72,6 +72,7 @@ void Window::WillCloseWindow(bool* prevent_default) {
 void Window::OnWindowClosed() {
   Emit("closed");
 
+  RemoveFromWeakMap();
   window_->RemoveObserver(this);
 }
 
@@ -542,7 +543,10 @@ void Initialize(v8::Local<v8::Object> exports, v8::Local<v8::Value> unused,
   v8::Local<v8::Function> constructor = mate::CreateConstructor<Window>(
       isolate, "BrowserWindow", base::Bind(&Window::New));
   mate::Dictionary browser_window(isolate, constructor);
-  browser_window.SetMethod("fromId", &mate::TrackableObject::FromWeakMapID);
+  browser_window.SetMethod("fromId",
+                           &mate::TrackableObject<Window>::FromWeakMapID);
+  browser_window.SetMethod("getAllWindows",
+                           &mate::TrackableObject<Window>::GetAll);
 
   mate::Dictionary dict(isolate, exports);
   dict.Set("BrowserWindow", browser_window);

atom/browser/api/atom_api_window.h

@@ -37,7 +37,7 @@ namespace api {
 
 class WebContents;
 
-class Window : public mate::TrackableObject,
+class Window : public mate::TrackableObject<Window>,
                public NativeWindowObserver {
  public:
   static mate::Wrappable* New(v8::Isolate* isolate,

atom/browser/api/lib/browser-window.coffee

@@ -42,10 +42,6 @@ BrowserWindow::setMenu = (menu) ->
   @menu = menu  # Keep a reference of menu in case of GC.
   @_setMenu menu
 
-BrowserWindow.getAllWindows = ->
-  windows = BrowserWindow.windows
-  windows.get key for key in windows.keys()
-
 BrowserWindow.getFocusedWindow = ->
   windows = BrowserWindow.getAllWindows()
   return window for window in windows when window.isFocused()

atom/browser/api/trackable_object.cc

@@ -4,6 +4,8 @@
 
 #include "atom/browser/api/trackable_object.h"
 
+#include "atom/browser/atom_browser_main_parts.h"
+#include "base/bind.h"
 #include "base/supports_user_data.h"
 
 namespace mate {
@@ -26,56 +28,41 @@ class IDUserData : public base::SupportsUserData::Data {
 
 }  // namespace
 
-atom::IDWeakMap TrackableObject::weak_map_;
+TrackableObjectBase::TrackableObjectBase()
-
+    : weak_map_id_(0), wrapped_(nullptr) {
-// static
-TrackableObject* TrackableObject::FromWeakMapID(v8::Isolate* isolate,
-                                                int32_t id) {
-  v8::MaybeLocal<v8::Object> object = weak_map_.Get(isolate, id);
-  if (object.IsEmpty())
-    return nullptr;
-
-  TrackableObject* self = nullptr;
-  mate::ConvertFromV8(isolate, object.ToLocalChecked(), &self);
-  return self;
 }
 
-// static
+TrackableObjectBase::~TrackableObjectBase() {
-TrackableObject* TrackableObject::FromWrappedClass(
-    v8::Isolate* isolate, base::SupportsUserData* wrapped) {
-  auto id = static_cast<IDUserData*>(wrapped->GetUserData(kTrackedObjectKey));
-  if (!id)
-    return nullptr;
-  return FromWeakMapID(isolate, *id);
 }
 
-// static
+void TrackableObjectBase::AfterInit(v8::Isolate* isolate) {
-void TrackableObject::ReleaseAllWeakReferences() {
-  weak_map_.Clear();
-}
-
-TrackableObject::TrackableObject() : weak_map_id_(0), wrapped_(nullptr) {
-}
-
-TrackableObject::~TrackableObject() {
-  weak_map_.Remove(weak_map_id_);
-}
-
-void TrackableObject::AfterInit(v8::Isolate* isolate) {
-  weak_map_id_ = weak_map_.Add(isolate, GetWrapper(isolate));
   if (wrapped_)
     AttachAsUserData(wrapped_);
 }
 
-void TrackableObject::AttachAsUserData(base::SupportsUserData* wrapped) {
+void TrackableObjectBase::AttachAsUserData(base::SupportsUserData* wrapped) {
   if (weak_map_id_ != 0) {
     wrapped->SetUserData(kTrackedObjectKey, new IDUserData(weak_map_id_));
     wrapped_ = nullptr;
   } else {
-    // If the TrackableObject is not ready yet then delay SetUserData until
+    // If the TrackableObjectBase is not ready yet then delay SetUserData until
     // AfterInit is called.
     wrapped_ = wrapped;
   }
 }
 
+// static
+int32_t TrackableObjectBase::GetIDFromWrappedClass(base::SupportsUserData* w) {
+  auto id = static_cast<IDUserData*>(w->GetUserData(kTrackedObjectKey));
+  if (id)
+    return *id;
+  else
+    return 0;
+}
+
+// static
+void TrackableObjectBase::RegisterDestructionCallback(void (*c)()) {
+  atom::AtomBrowserMainParts::Get()->RegisterDestructionCallback(base::Bind(c));
+}
+
 }  // namespace mate

atom/browser/api/trackable_object.h

@@ -5,6 +5,8 @@
 #ifndef ATOM_BROWSER_API_TRACKABLE_OBJECT_H_
 #define ATOM_BROWSER_API_TRACKABLE_OBJECT_H_
 
+#include <vector>
+
 #include "atom/browser/api/event_emitter.h"
 #include "atom/common/id_weak_map.h"
 
@@ -14,42 +16,101 @@ class SupportsUserData;
 
 namespace mate {
 
-// All instances of TrackableObject will be kept in a weak map and can be got
+// Users should use TrackableObject instead.
-// from its ID.
+class TrackableObjectBase : public mate::EventEmitter {
-class TrackableObject : public mate::EventEmitter {
  public:
-  // Find out the TrackableObject from its ID in weak map.
+  TrackableObjectBase();
-  static TrackableObject* FromWeakMapID(v8::Isolate* isolate, int32_t id);
-
-  // Find out the TrackableObject from the class it wraps.
-  static TrackableObject* FromWrappedClass(
-      v8::Isolate* isolate, base::SupportsUserData* wrapped);
-
-  // Releases all weak references in weak map, called when app is terminating.
-  static void ReleaseAllWeakReferences();
-
-  // mate::Wrappable:
-  void AfterInit(v8::Isolate* isolate) override;
 
   // The ID in weak map.
   int32_t weak_map_id() const { return weak_map_id_; }
 
- protected:
-  TrackableObject();
-  ~TrackableObject() override;
-
   // Wrap TrackableObject into a class that SupportsUserData.
   void AttachAsUserData(base::SupportsUserData* wrapped);
 
- private:
+ protected:
-  static atom::IDWeakMap weak_map_;
+  ~TrackableObjectBase() override;
+
+  // mate::Wrappable:
+  void AfterInit(v8::Isolate* isolate) override;
+
+  // Get the weak_map_id from SupportsUserData.
+  static int32_t GetIDFromWrappedClass(base::SupportsUserData* wrapped);
+
+  // Register a callback that should be destroyed before JavaScript environment
+  // gets destroyed.
+  static void RegisterDestructionCallback(void (*callback)());
 
   int32_t weak_map_id_;
   base::SupportsUserData* wrapped_;
 
+ private:
+  DISALLOW_COPY_AND_ASSIGN(TrackableObjectBase);
+};
+
+// All instances of TrackableObject will be kept in a weak map and can be got
+// from its ID.
+template<typename T>
+class TrackableObject : public TrackableObjectBase {
+ public:
+  // Finds out the TrackableObject from its ID in weak map.
+  static T* FromWeakMapID(v8::Isolate* isolate, int32_t id) {
+    v8::MaybeLocal<v8::Object> object = weak_map_.Get(isolate, id);
+    if (object.IsEmpty())
+      return nullptr;
+
+    T* self = nullptr;
+    mate::ConvertFromV8(isolate, object.ToLocalChecked(), &self);
+    return self;
+  }
+
+  // Finds out the TrackableObject from the class it wraps.
+  static T* FromWrappedClass(v8::Isolate* isolate,
+                             base::SupportsUserData* wrapped) {
+    int32_t id = GetIDFromWrappedClass(wrapped);
+    if (!id)
+      return nullptr;
+    return FromWeakMapID(isolate, id);
+  }
+
+  // Returns all objects in this class's weak map.
+  static std::vector<v8::Local<v8::Object>> GetAll(v8::Isolate* isolate) {
+    return weak_map_.Values(isolate);
+  }
+
+  TrackableObject() {
+    RegisterDestructionCallback(&TrackableObject<T>::ReleaseAllWeakReferences);
+  }
+
+  // Removes this instance from the weak map.
+  void RemoveFromWeakMap() {
+    if (weak_map_.Has(weak_map_id()))
+      weak_map_.Remove(weak_map_id());
+  }
+
+ protected:
+  ~TrackableObject() override {
+    RemoveFromWeakMap();
+  }
+
+  void AfterInit(v8::Isolate* isolate) override {
+    weak_map_id_ = weak_map_.Add(isolate, GetWrapper(isolate));
+    TrackableObjectBase::AfterInit(isolate);
+  }
+
+ private:
+  // Releases all weak references in weak map, called when app is terminating.
+  static void ReleaseAllWeakReferences() {
+    weak_map_.Clear();
+  }
+
+  static atom::IDWeakMap weak_map_;
+
   DISALLOW_COPY_AND_ASSIGN(TrackableObject);
 };
 
+template<typename T>
+atom::IDWeakMap TrackableObject<T>::weak_map_;
+
 }  // namespace mate
 
 #endif  // ATOM_BROWSER_API_TRACKABLE_OBJECT_H_

atom/browser/atom_browser_main_parts.cc

@@ -37,7 +37,8 @@ AtomBrowserMainParts::AtomBrowserMainParts()
 }
 
 AtomBrowserMainParts::~AtomBrowserMainParts() {
-  mate::TrackableObject::ReleaseAllWeakReferences();
+  for (const auto& callback : destruction_callbacks_)
+    callback.Run();
 }
 
 // static
@@ -46,6 +47,11 @@ AtomBrowserMainParts* AtomBrowserMainParts::Get() {
   return self_;
 }
 
+void AtomBrowserMainParts::RegisterDestructionCallback(
+    const base::Closure& callback) {
+  destruction_callbacks_.push_back(callback);
+}
+
 brightray::BrowserContext* AtomBrowserMainParts::CreateBrowserContext() {
   return new AtomBrowserContext();
 }

atom/browser/atom_browser_main_parts.h

@@ -5,6 +5,9 @@
 #ifndef ATOM_BROWSER_ATOM_BROWSER_MAIN_PARTS_H_
 #define ATOM_BROWSER_ATOM_BROWSER_MAIN_PARTS_H_
 
+#include <list>
+
+#include "base/callback.h"
 #include "base/timer/timer.h"
 #include "brightray/browser/browser_main_parts.h"
 
@@ -24,6 +27,10 @@ class AtomBrowserMainParts : public brightray::BrowserMainParts {
 
   static AtomBrowserMainParts* Get();
 
+  // Register a callback that should be destroyed before JavaScript environment
+  // gets destroyed.
+  void RegisterDestructionCallback(const base::Closure& callback);
+
   Browser* browser() { return browser_.get(); }
 
  protected:
@@ -53,6 +60,9 @@ class AtomBrowserMainParts : public brightray::BrowserMainParts {
 
   base::Timer gc_timer_;
 
+  // List of callbacks should be executed before destroying JS env.
+  std::list<base::Closure> destruction_callbacks_;
+
   static AtomBrowserMainParts* self_;
 
   DISALLOW_COPY_AND_ASSIGN(AtomBrowserMainParts);