Remove itself from the cleanup list when it is destroyed

atom/browser/api/trackable_object.cc

@@ -30,10 +30,11 @@ class IDUserData : public base::SupportsUserData::Data {
 
 TrackableObjectBase::TrackableObjectBase()
     : weak_map_id_(0), wrapped_(nullptr), weak_factory_(this) {
-  RegisterDestructionCallback(GetDestroyClosure());
+  cleanup_ = RegisterDestructionCallback(GetDestroyClosure());
 }
 
 TrackableObjectBase::~TrackableObjectBase() {
+  cleanup_.Run();
 }
 
 void TrackableObjectBase::AfterInit(v8::Isolate* isolate) {
@@ -74,9 +75,9 @@ int32_t TrackableObjectBase::GetIDFromWrappedClass(base::SupportsUserData* w) {
 }
 
 // static
-void TrackableObjectBase::RegisterDestructionCallback(
-    const base::Closure& closure) {
-  atom::AtomBrowserMainParts::Get()->RegisterDestructionCallback(closure);
+base::Closure TrackableObjectBase::RegisterDestructionCallback(
+    const base::Closure& c) {
+  return atom::AtomBrowserMainParts::Get()->RegisterDestructionCallback(c);
 }
 
 }  // namespace mate

atom/browser/api/trackable_object.h

@@ -48,7 +48,7 @@ class TrackableObjectBase : public mate::EventEmitter {
 
   // Register a callback that should be destroyed before JavaScript environment
   // gets destroyed.
-  static void RegisterDestructionCallback(const base::Closure& closure);
+  static base::Closure RegisterDestructionCallback(const base::Closure& c);
 
   int32_t weak_map_id_;
   base::SupportsUserData* wrapped_;
@@ -56,6 +56,7 @@ class TrackableObjectBase : public mate::EventEmitter {
  private:
   void Destroy();
 
+  base::Closure cleanup_;
   base::WeakPtrFactory<TrackableObjectBase> weak_factory_;
 
   DISALLOW_COPY_AND_ASSIGN(TrackableObjectBase);

atom/browser/atom_browser_main_parts.cc

@@ -25,6 +25,11 @@
 
 namespace atom {
 
+template<typename T>
+void Erase(T* container, typename T::iterator iter) {
+  container->erase(iter);
+}
+
 // static
 AtomBrowserMainParts* AtomBrowserMainParts::self_ = NULL;
 
@@ -56,9 +61,10 @@ bool AtomBrowserMainParts::SetExitCode(int code) {
   return true;
 }
 
-void AtomBrowserMainParts::RegisterDestructionCallback(
+base::Closure AtomBrowserMainParts::RegisterDestructionCallback(
     const base::Closure& callback) {
-  destruction_callbacks_.push_back(callback);
+  auto iter = destructors_.insert(destructors_.end(), callback);
+  return base::Bind(&Erase<std::list<base::Closure>>, &destructors_, iter);
 }
 
 void AtomBrowserMainParts::PreEarlyInitialization() {
@@ -150,8 +156,13 @@ void AtomBrowserMainParts::PostMainMessageLoopRun() {
   // Make sure destruction callbacks are called before message loop is
   // destroyed, otherwise some objects that need to be deleted on IO thread
   // won't be freed.
-  for (const auto& callback : destruction_callbacks_)
+  // We don't use ranged for loop because iterators are getting invalided when
+  // the callback runs.
+  for (auto iter = destructors_.begin(); iter != destructors_.end();) {
+    base::Closure& callback = *iter;
+    ++iter;
     callback.Run();
+  }
 
   // Destroy JavaScript environment immediately after running destruction
   // callbacks.

atom/browser/atom_browser_main_parts.h

@@ -36,7 +36,8 @@ class AtomBrowserMainParts : public brightray::BrowserMainParts {
 
   // Register a callback that should be destroyed before JavaScript environment
   // gets destroyed.
-  void RegisterDestructionCallback(const base::Closure& callback);
+  // Returns a closure that can be used to remove |callback| from the list.
+  base::Closure RegisterDestructionCallback(const base::Closure& callback);
 
   Browser* browser() { return browser_.get(); }
 
@@ -82,7 +83,7 @@ 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_;
+  std::list<base::Closure> destructors_;
 
   static AtomBrowserMainParts* self_;