Make Session trackable

atom/browser/api/atom_api_app.cc

@@ -176,7 +176,7 @@ v8::Local<v8::Value> App::DefaultSession(v8::Isolate* isolate) {
   if (default_session_.IsEmpty()) {
     auto browser_context = static_cast<AtomBrowserContext*>(
         AtomBrowserMainParts::Get()->browser_context());
-    auto handle = Session::Create(isolate, browser_context);
+    auto handle = Session::CreateFrom(isolate, browser_context);
     default_session_.Reset(isolate, handle.ToV8());
   }
   return v8::Local<v8::Value>::New(isolate, default_session_);

atom/browser/api/atom_api_session.cc

@@ -68,6 +68,7 @@ class ResolveProxyHelper {
 
 Session::Session(AtomBrowserContext* browser_context)
     : browser_context_(browser_context) {
+  AttachAsUserData(browser_context);
 }
 
 Session::~Session() {
@@ -93,9 +94,13 @@ mate::ObjectTemplateBuilder Session::GetObjectTemplateBuilder(
 }
 
 // static
-mate::Handle<Session> Session::Create(
+mate::Handle<Session> Session::CreateFrom(
     v8::Isolate* isolate,
     AtomBrowserContext* browser_context) {
+  auto existing = TrackableObject::FromWrappedClass(isolate, browser_context);
+  if (existing)
+    return mate::CreateHandle(isolate, static_cast<Session*>(existing));
+
   return mate::CreateHandle(isolate, new Session(browser_context));
 }
 

atom/browser/api/atom_api_session.h

@@ -7,9 +7,9 @@
 
 #include <string>
 
+#include "atom/browser/api/trackable_object.h"
 #include "base/callback.h"
 #include "native_mate/handle.h"
-#include "native_mate/wrappable.h"
 
 class GURL;
 
@@ -19,12 +19,13 @@ class AtomBrowserContext;
 
 namespace api {
 
-class Session: public mate::Wrappable {
+class Session: public mate::TrackableObject {
  public:
   using ResolveProxyCallback = base::Callback<void(std::string)>;
 
-  static mate::Handle<Session> Create(v8::Isolate* isolate,
-                                      AtomBrowserContext* browser_context);
+  // Gets or creates Session from the |browser_context|.
+  static mate::Handle<Session> CreateFrom(
+      v8::Isolate* isolate, AtomBrowserContext* browser_context);
 
  protected:
   explicit Session(AtomBrowserContext* browser_context);

atom/browser/api/atom_api_web_contents.cc

@@ -155,6 +155,7 @@ WebContents::WebContents(content::WebContents* web_contents)
       auto_size_enabled_(false),
       is_full_page_plugin_(false),
       inspectable_web_contents_(nullptr) {
+  AttachAsUserData(web_contents);
 }
 
 WebContents::WebContents(const mate::Dictionary& options)
@@ -176,6 +177,7 @@ WebContents::WebContents(const mate::Dictionary& options)
   if (options.Get("embedder", &embedder) && embedder)
     owner_window = NativeWindow::FromWebContents(embedder->web_contents());
 
+  AttachAsUserData(web_contents);
   InitWithWebContents(web_contents, owner_window);
   inspectable_web_contents_ = managed_web_contents();
 
@@ -608,7 +610,7 @@ void WebContents::InspectServiceWorker() {
 
 v8::Local<v8::Value> WebContents::Session(v8::Isolate* isolate) {
   if (session_.IsEmpty()) {
-    mate::Handle<api::Session> handle = Session::Create(
+    mate::Handle<api::Session> handle = Session::CreateFrom(
         isolate,
         static_cast<AtomBrowserContext*>(web_contents()->GetBrowserContext()));
     session_.Reset(isolate, handle.ToV8());
@@ -784,11 +786,6 @@ bool WebContents::IsGuest() const {
   return is_guest();
 }
 
-void WebContents::AfterInit(v8::Isolate* isolate) {
-  mate::TrackableObject::AfterInit(isolate);
-  AttachAsUserData(web_contents());
-}
-
 mate::ObjectTemplateBuilder WebContents::GetObjectTemplateBuilder(
     v8::Isolate* isolate) {
   if (template_.IsEmpty())

atom/browser/api/atom_api_web_contents.h

@@ -142,7 +142,6 @@ class WebContents : public mate::TrackableObject,
   ~WebContents();
 
   // mate::Wrappable:
-  void AfterInit(v8::Isolate* isolate);
   mate::ObjectTemplateBuilder GetObjectTemplateBuilder(
       v8::Isolate* isolate) override;
 
@@ -241,7 +240,6 @@ class WebContents : public mate::TrackableObject,
   // Returns the default size of the guestview.
   gfx::Size GetDefaultSize() const;
 
-
   v8::Global<v8::Value> session_;
 
   // Stores whether the contents of the guest can be transparent.

atom/browser/api/trackable_object.cc

@@ -54,7 +54,7 @@ void TrackableObject::ReleaseAllWeakReferences() {
   weak_map_.Clear();
 }
 
-TrackableObject::TrackableObject() : weak_map_id_(0) {
+TrackableObject::TrackableObject() : weak_map_id_(0), wrapped_(nullptr) {
 }
 
 TrackableObject::~TrackableObject() {
@@ -63,10 +63,19 @@ TrackableObject::~TrackableObject() {
 
 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) {
-  wrapped->SetUserData(kTrackedObjectKey, new IDUserData(weak_map_id_));
+  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
+    // AfterInit is called.
+    wrapped_ = wrapped;
+  }
 }
 
 }  // namespace mate

atom/browser/api/trackable_object.h

@@ -45,6 +45,7 @@ class TrackableObject : public mate::EventEmitter {
   static atom::IDWeakMap weak_map_;
 
   int32_t weak_map_id_;
+  base::SupportsUserData* wrapped_;
 
   DISALLOW_COPY_AND_ASSIGN(TrackableObject);
 };