framework: fix race condition in shceduler

framework/cybertron/croutine/croutine.cpp

@@ -57,8 +57,8 @@ RoutineState CRoutine::Resume() {
     return RoutineState::FINISHED;
   }
 
+  std::unique_lock<std::mutex> ul(op_mtx_);
   UpdateState();
-
   // Keep compatibility with different policies.
   if (!IsRunning() && !IsReady()) {
     if (IsWaitingInput()) {

framework/cybertron/croutine/croutine.h

@@ -75,18 +75,9 @@ class CRoutine {
     SetState(RoutineState::FINISHED);
   }
 
-  inline bool TryLockForOp() {
-    std::lock_guard<std::mutex> lh(op_mtx_);
-    if (!being_op_) {
-      being_op_ = true;
-      return being_op_;
-    }
-    return false;
-  }
-
-  inline void TryUnlockForOp() {
-    std::lock_guard<std::mutex> lh(op_mtx_);
-    being_op_ = false;
+  inline std::unique_lock<std::mutex> GetLock() const {
+    std::unique_lock<std::mutex> ul(op_mtx_, std::defer_lock);
+    return ul;
   }
 
   inline void SetState(const RoutineState &state, bool is_notify = false) {
@@ -205,8 +196,7 @@ class CRoutine {
   double vruntime_ = 0.0;
   double normalized_vruntime_ = 0.0;
   bool force_stop_ = false;
-  bool being_op_ = false;
-  std::mutex op_mtx_;
+  mutable std::mutex op_mtx_;
 };
 
 }  // namespace croutine

framework/cybertron/scheduler/policy/cfs.cpp

@@ -48,6 +48,7 @@ std::shared_ptr<CRoutine> CFSContext::NextRoutine() {
   }
   return cr;
 }
+
 std::shared_ptr<CRoutine> CFSContext::NextLocalRoutine() {
   std::lock_guard<std::mutex> lock(mtx_run_queue_);
   auto start_perf_time = apollo::cybertron::Time::Now().ToNanosecond();
@@ -67,16 +68,20 @@ std::shared_ptr<CRoutine> CFSContext::NextLocalRoutine() {
   std::shared_ptr<CRoutine> croutine = nullptr;
   for (auto it = local_rb_map_.begin(); it != local_rb_map_.end();) {
     auto cr = it->second;
-    if (!cr->TryLockForOp()) {
+    auto lock = cr->GetLock();
+    if (!lock.try_lock()) {
       ++it;
       continue;
     }
 
-    auto cr_id = cr->Id();
     cr->UpdateState();
+    if (cr->IsRunning()) {
+      ++it;
+      continue;
+    }
+
     if (cr->IsFinished()) {
       it = local_rb_map_.erase(it);
-      cr->TryUnlockForOp();
       continue;
     }
 
@@ -86,10 +91,8 @@ std::shared_ptr<CRoutine> CFSContext::NextLocalRoutine() {
       croutine->SetState(RoutineState::RUNNING);
       cur_croutine_ = croutine;
       local_rb_map_.erase(it);
-      cr->TryUnlockForOp();
       break;
     }
-    cr->TryUnlockForOp();
     ++it;
   }
   if (croutine == nullptr) {
@@ -112,7 +115,8 @@ std::shared_ptr<CRoutine> CFSContext::NextAffinityRoutine() {
 
   for (auto it = affinity_rb_map_.begin(); it != affinity_rb_map_.end();) {
     auto cr = it->second;
-    if (!cr->TryLockForOp()) {
+    auto lock = cr->GetLock();
+    if (!lock.try_lock()) {
       ++it;
       continue;
     }
@@ -120,17 +124,19 @@ std::shared_ptr<CRoutine> CFSContext::NextAffinityRoutine() {
     cr->UpdateState();
     if (cr->IsFinished()) {
       it = affinity_rb_map_.erase(it);
-      cr->TryUnlockForOp();
+      continue;
+    }
+
+    if (cr->IsRunning()) {
+      ++it;
       continue;
     }
 
     if (cr->IsReady()) {
       croutine = cr;
       cr->SetState(RoutineState::RUNNING);
-      cr->TryUnlockForOp();
       break;
     }
-    cr->TryUnlockForOp();
     ++it;
   }
   if (croutine) {

framework/cybertron/scheduler/policy/classic.cpp

@@ -48,7 +48,8 @@ std::shared_ptr<CRoutine> ClassicContext::NextRoutine() {
   std::shared_ptr<CRoutine> croutine = nullptr;
   for (auto it = rb_map_.begin(); it != rb_map_.end();) {
     auto cr = it->second;
-    if (!cr->TryLockForOp()) {
+    auto lock = cr->GetLock();
+    if (!lock.try_lock()) {
       ++it;
       continue;
     }
@@ -58,17 +59,14 @@ std::shared_ptr<CRoutine> ClassicContext::NextRoutine() {
 
     if (cr->IsFinished()) {
       it = rb_map_.erase(it);
-      cr->TryUnlockForOp();
       continue;
     }
 
     if (cr->IsReady()) {
       croutine = cr;
       cr->SetState(RoutineState::RUNNING);
-      cr->TryUnlockForOp();
       break;
     }
-    cr->TryUnlockForOp();
     ++it;
   }
   if (croutine == nullptr) {