From 6ec6a72661f303de30032c5de9ba09af2a97948e Mon Sep 17 00:00:00 2001
From: ziyuehuang <ziyuehuang@tencent.com>
Date: Wed, 11 Jan 2023 13:32:15 +0800
Subject: [PATCH] fix thread safety of RunCallback

---
 include/ps/kv_app.h | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/include/ps/kv_app.h b/include/ps/kv_app.h
index 50bdde6b..914112d8 100644
--- a/include/ps/kv_app.h
+++ b/include/ps/kv_app.h
@@ -640,18 +640,16 @@ void KVWorker<Val>::Process(const Message& msg) {
 }
 template <typename Val>
 void KVWorker<Val>::RunCallback(int timestamp) {
+  Callback cb = nullptr;
   mu_.lock();
   auto it = callbacks_.find(timestamp);
   if (it != callbacks_.end()) {
-    mu_.unlock();
-
-    CHECK(it->second);
-    it->second();
-
-    mu_.lock();
+    cb = it->second;
+    CHECK(cb);
     callbacks_.erase(it);
   }
   mu_.unlock();
+  if (cb) cb();
 }
 
 template <typename Val>