From 375dce4529ed183e3eba5bf349f1298a9fd6c8ca Mon Sep 17 00:00:00 2001
From: kswang <wangkaisheng2@huawei.com>
Date: Wed, 30 Dec 2020 22:26:18 +0800
Subject: [PATCH] fix exception segmentfault

---
 mindspore/core/ir/tensor.h          |  3 +--
 mindspore/core/utils/ms_exception.h | 17 ++++++-----------
 2 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/mindspore/core/ir/tensor.h b/mindspore/core/ir/tensor.h
index b2eadd4d1a..f68df3cd5b 100644
--- a/mindspore/core/ir/tensor.h
+++ b/mindspore/core/ir/tensor.h
@@ -87,10 +87,9 @@ class WaitEvent : public ExceptionListener {
     if (!need_wait_) {
       return;
     }
-    MsException::Instance().AddExceptionListener(const_cast<WaitEvent *>(this));
+    MsException::Instance().SetExceptionListener(const_cast<WaitEvent *>(this));
     cond_var_.wait(lock, [this] { return !need_wait_; });
     MsException::Instance().CheckException();
-    MsException::Instance().RemoveExceptionListener(const_cast<WaitEvent *>(this));
   }
 
   void set_need_wait(bool need_wait) {
diff --git a/mindspore/core/utils/ms_exception.h b/mindspore/core/utils/ms_exception.h
index 8cc013fb47..0e8462b1da 100644
--- a/mindspore/core/utils/ms_exception.h
+++ b/mindspore/core/utils/ms_exception.h
@@ -34,13 +34,10 @@ class MsException {
 
   void SetException() {
     exception_ptr_ = std::current_exception();
-    if (exception_ptr_ != nullptr) {
-      for (auto &listener : listeners_) {
-        if (listener == nullptr) {
-          continue;
-        }
-        listener->OnException();
-      }
+    if (exception_ptr_ != nullptr && listener_ != nullptr) {
+      auto listener = listener_;
+      listener_ = nullptr;
+      listener->OnException();
     }
   }
 
@@ -52,15 +49,13 @@ class MsException {
     }
   }
 
-  void AddExceptionListener(ExceptionListener *listener) { (void)listeners_.insert(listener); }
-
-  void RemoveExceptionListener(ExceptionListener *listener) { (void)listeners_.erase(listener); }
+  void SetExceptionListener(ExceptionListener *listener) { listener_ = listener; }
 
  private:
   MsException() = default;
   ~MsException() = default;
   DISABLE_COPY_AND_ASSIGN(MsException)
-  std::set<ExceptionListener *> listeners_;
+  ExceptionListener *listener_{nullptr};
   std::exception_ptr exception_ptr_{nullptr};
 };
 }  // namespace mindspore