From 69d6f6c347f6f492c29c0158c5c157eae608c0d5 Mon Sep 17 00:00:00 2001
From: wxyu <xy.wang@zilliz.com>
Date: Sat, 12 Oct 2019 16:49:06 +0800
Subject: [PATCH] MS-641 Segment fault(signal 11) in PickToLoad

Former-commit-id: 9e97782faf0f24cfc89ca36e05e037f82489c8cd
---
 cpp/CHANGELOG.md                | 1 +
 cpp/src/scheduler/TaskTable.cpp | 9 ++++++---
 cpp/src/scheduler/TaskTable.h   | 5 +----
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/cpp/CHANGELOG.md b/cpp/CHANGELOG.md
index d19723fede..8d25c19ce8 100644
--- a/cpp/CHANGELOG.md
+++ b/cpp/CHANGELOG.md
@@ -15,6 +15,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-620 - Get table row counts display wrong error code
 - MS-637 - out of memory when load too many tasks
 - MS-640 - Cache object size calculate incorrect
+- MS-641 - Segment fault(signal 11) in PickToLoad
 
 ## Improvement
 - MS-552 - Add and change the easylogging library
diff --git a/cpp/src/scheduler/TaskTable.cpp b/cpp/src/scheduler/TaskTable.cpp
index d03d36359e..2f7576de34 100644
--- a/cpp/src/scheduler/TaskTable.cpp
+++ b/cpp/src/scheduler/TaskTable.cpp
@@ -158,8 +158,9 @@ TaskTableItem::Dump() {
 
 std::vector<uint64_t>
 TaskTable::PickToLoad(uint64_t limit) {
+    std::lock_guard<std::mutex> lock(mutex_);
     size_t count = 0;
-    for (int j = last_finish_ + 1; j < table_.size(); ++j) {
+    for (uint64_t j = last_finish_ + 1; j < table_.size(); ++j) {
         if (not table_[j]) {
             SERVER_LOG_WARNING << "table[" << j << "] is nullptr";
         }
@@ -186,6 +187,7 @@ TaskTable::PickToLoad(uint64_t limit) {
 
 std::vector<uint64_t>
 TaskTable::PickToExecute(uint64_t limit) {
+    std::lock_guard<std::mutex> lock(mutex_);
     std::vector<uint64_t> indexes;
     bool cross = false;
     for (uint64_t i = last_finish_ + 1, count = 0; i < table_.size() && count < limit; ++i) {
@@ -202,7 +204,7 @@ TaskTable::PickToExecute(uint64_t limit) {
 
 void
 TaskTable::Put(TaskPtr task) {
-    std::lock_guard<std::mutex> lock(id_mutex_);
+    std::lock_guard<std::mutex> lock(mutex_);
     auto item = std::make_shared<TaskTableItem>();
     item->id = id_++;
     item->task = std::move(task);
@@ -216,7 +218,7 @@ TaskTable::Put(TaskPtr task) {
 
 void
 TaskTable::Put(std::vector<TaskPtr>& tasks) {
-    std::lock_guard<std::mutex> lock(id_mutex_);
+    std::lock_guard<std::mutex> lock(mutex_);
     for (auto& task : tasks) {
         auto item = std::make_shared<TaskTableItem>();
         item->id = id_++;
@@ -232,6 +234,7 @@ TaskTable::Put(std::vector<TaskPtr>& tasks) {
 
 TaskTableItemPtr
 TaskTable::Get(uint64_t index) {
+    std::lock_guard<std::mutex> lock(mutex_);
     return table_[index];
 }
 
diff --git a/cpp/src/scheduler/TaskTable.h b/cpp/src/scheduler/TaskTable.h
index ad81b5d439..39e129ab0e 100644
--- a/cpp/src/scheduler/TaskTable.h
+++ b/cpp/src/scheduler/TaskTable.h
@@ -149,9 +149,6 @@ class TaskTable {
     }
 
  public:
-    TaskTableItemPtr& operator[](uint64_t index) {
-        return table_[index];
-    }
 
     std::deque<TaskTableItemPtr>::iterator
     begin() {
@@ -244,7 +241,7 @@ class TaskTable {
 
  private:
     std::uint64_t id_ = 0;
-    mutable std::mutex id_mutex_;
+    mutable std::mutex mutex_;
     std::deque<TaskTableItemPtr> table_;
     std::function<void(void)> subscriber_ = nullptr;