From 35e28f81b5a14029ae7ae8e9b5bb494ff7c9dd30 Mon Sep 17 00:00:00 2001
From: wxyu <xy.wang@zilliz.com>
Date: Sat, 24 Aug 2019 16:36:56 +0800
Subject: [PATCH] MS-414 Add TaskType in Scheduler::Task

Former-commit-id: d470511e0d22832db09607ad017c6dd4a697cf05
---
 cpp/CHANGELOG.md                      |  1 +
 cpp/src/scheduler/task/DeleteTask.cpp |  2 +-
 cpp/src/scheduler/task/SearchTask.cpp |  5 +++--
 cpp/src/scheduler/task/Task.h         | 13 ++++++++++++-
 cpp/src/scheduler/task/TestTask.cpp   |  2 ++
 cpp/src/scheduler/task/TestTask.h     |  2 +-
 6 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/cpp/CHANGELOG.md b/cpp/CHANGELOG.md
index 671d865a..3b4888de 100644
--- a/cpp/CHANGELOG.md
+++ b/cpp/CHANGELOG.md
@@ -50,6 +50,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-408 - Add device_id in resource construct function
 - MS-409 - Using new scheduler
 - MS-410 - Add resource config comment
+- MS-414 - Add TaskType in Scheduler::Task
 
 ## New Feature
 - MS-343 - Implement ResourceMgr
diff --git a/cpp/src/scheduler/task/DeleteTask.cpp b/cpp/src/scheduler/task/DeleteTask.cpp
index 9f76de50..0e9c7dc3 100644
--- a/cpp/src/scheduler/task/DeleteTask.cpp
+++ b/cpp/src/scheduler/task/DeleteTask.cpp
@@ -12,7 +12,7 @@ namespace milvus {
 namespace engine {
 
 XDeleteTask::XDeleteTask(DeleteContextPtr &delete_context)
-    : delete_context_ptr_(delete_context) {}
+    : Task(TaskType::DeleteTask), delete_context_ptr_(delete_context) {}
 
 void
 XDeleteTask::Load(LoadType type, uint8_t device_id) {
diff --git a/cpp/src/scheduler/task/SearchTask.cpp b/cpp/src/scheduler/task/SearchTask.cpp
index 7d342a35..3ddc2beb 100644
--- a/cpp/src/scheduler/task/SearchTask.cpp
+++ b/cpp/src/scheduler/task/SearchTask.cpp
@@ -81,11 +81,12 @@ CollectFileMetrics(int file_type, size_t file_size) {
     }
 }
 
-XSearchTask::XSearchTask(TableFileSchemaPtr file) : file_(file) {
+XSearchTask::XSearchTask(TableFileSchemaPtr file)
+    : Task(TaskType::SearchTask), file_(file) {
     index_engine_ = EngineFactory::Build(file_->dimension_,
                                          file_->location_,
                                          (EngineType) file_->engine_type_,
-                                         (MetricType)file_->metric_type_,
+                                         (MetricType) file_->metric_type_,
                                          file_->nlist_);
 }
 
diff --git a/cpp/src/scheduler/task/Task.h b/cpp/src/scheduler/task/Task.h
index 4d78894d..bc2dae46 100644
--- a/cpp/src/scheduler/task/Task.h
+++ b/cpp/src/scheduler/task/Task.h
@@ -21,13 +21,20 @@ enum class LoadType {
     GPU2CPU,
 };
 
+enum class TaskType {
+    SearchTask,
+    DeleteTask,
+    TestTask,
+};
+
 class Task;
 
 using TaskPtr = std::shared_ptr<Task>;
 
 class Task {
 public:
-    Task() = default;
+    explicit
+    Task(TaskType type) : type_(type) {}
 
     virtual void
     Load(LoadType type, uint8_t device_id) = 0;
@@ -39,9 +46,13 @@ public:
     virtual TaskPtr
     Clone() = 0;
 
+    inline TaskType
+    Type() const { return type_; }
+
 public:
     std::vector<SearchContextPtr> search_contexts_;
     ScheduleTaskPtr task_;
+    TaskType type_;
 };
 
 
diff --git a/cpp/src/scheduler/task/TestTask.cpp b/cpp/src/scheduler/task/TestTask.cpp
index 527e1b0f..29f2a940 100644
--- a/cpp/src/scheduler/task/TestTask.cpp
+++ b/cpp/src/scheduler/task/TestTask.cpp
@@ -11,6 +11,8 @@ namespace zilliz {
 namespace milvus {
 namespace engine {
 
+TestTask::TestTask() : Task(TaskType::TestTask) {}
+
 void
 TestTask::Load(LoadType type, uint8_t device_id) {
     load_count_++;
diff --git a/cpp/src/scheduler/task/TestTask.h b/cpp/src/scheduler/task/TestTask.h
index 0eb77504..5a466ec3 100644
--- a/cpp/src/scheduler/task/TestTask.h
+++ b/cpp/src/scheduler/task/TestTask.h
@@ -14,7 +14,7 @@ namespace engine {
 
 class TestTask : public Task {
 public:
-    TestTask() = default;
+    TestTask();
 
 public:
     void
-- 
GitLab