Exit the process gracefully

be/src/agent/agent_server.cpp

@@ -96,7 +96,33 @@ AgentServer::AgentServer(ExecEnv* exec_env, const TMasterInfo& master_info)
 #undef CREATE_AND_START_POOL
 }
 
-AgentServer::~AgentServer() = default;
+AgentServer::~AgentServer() {
+#ifndef STOP_POOL
+#define STOP_POOL(type, pool_name) pool_name->stop();
+#endif
+
+    STOP_POOL(CREATE_TABLE, _create_tablet_workers);
+    STOP_POOL(DROP_TABLE, _drop_tablet_workers);
+    // Both PUSH and REALTIME_PUSH type use _push_workers
+    STOP_POOL(PUSH, _push_workers);
+    STOP_POOL(PUBLISH_VERSION, _publish_version_workers);
+    STOP_POOL(CLEAR_TRANSACTION_TASK, _clear_transaction_task_workers);
+    STOP_POOL(DELETE, _delete_workers);
+    STOP_POOL(ALTER_TABLE, _alter_tablet_workers);
+    STOP_POOL(CLONE, _clone_workers);
+    STOP_POOL(STORAGE_MEDIUM_MIGRATE, _storage_medium_migrate_workers);
+    STOP_POOL(CHECK_CONSISTENCY, _check_consistency_workers);
+    STOP_POOL(REPORT_TASK, _report_task_workers);
+    STOP_POOL(REPORT_DISK_STATE, _report_disk_state_workers);
+    STOP_POOL(REPORT_OLAP_TABLE, _report_tablet_workers);
+    STOP_POOL(UPLOAD, _upload_workers);
+    STOP_POOL(DOWNLOAD, _download_workers);
+    STOP_POOL(MAKE_SNAPSHOT, _make_snapshot_workers);
+    STOP_POOL(RELEASE_SNAPSHOT, _release_snapshot_workers);
+    STOP_POOL(MOVE, _move_dir_workers);
+    STOP_POOL(UPDATE_TABLET_META_INFO, _update_tablet_meta_info_workers);
+#undef STOP_POOL
+}
 
 // TODO(lingbin): each task in the batch may have it own status or FE must check and
 // resend request when something is wrong(BE may need some logic to guarantee idempotence.

be/src/agent/agent_server.h

@@ -31,6 +31,7 @@
 
 namespace starrocks {
 
+class MultiWorkerPool;
 class TaskWorkerPool;
 
 // Each method corresponds to one RPC from FE Master, see BackendService.
@@ -62,7 +63,7 @@ class AgentServer {
     std::unique_ptr<TaskWorkerPool> _create_tablet_workers;
     std::unique_ptr<TaskWorkerPool> _drop_tablet_workers;
     std::unique_ptr<TaskWorkerPool> _push_workers;
-    std::unique_ptr<TaskWorkerPool> _publish_version_workers;
+    std::unique_ptr<MultiWorkerPool> _publish_version_workers;
     std::unique_ptr<TaskWorkerPool> _clear_transaction_task_workers;
     std::unique_ptr<TaskWorkerPool> _delete_workers;
     std::unique_ptr<TaskWorkerPool> _alter_tablet_workers;

be/src/agent/multi_worker_pool.cpp

@@ -4,9 +4,8 @@
 
 namespace starrocks {
 
-MultiWorkerPool::MultiWorkerPool(const TaskWorkerType worker_type, ExecEnv* env, const TMasterInfo& master_info,
-                                 int worker_num)
-        : TaskWorkerPool(worker_type, env, master_info, worker_num) {
+MultiWorkerPool::MultiWorkerPool(const TaskWorkerPool::TaskWorkerType worker_type, ExecEnv* env,
+                                 const TMasterInfo& master_info, int worker_num) {
     DCHECK(worker_num > 0);
     for (int i = 0; i < worker_num; i++) {
         auto pool = std::make_shared<TaskWorkerPool>(worker_type, env, master_info, 1);
@@ -20,6 +19,12 @@ void MultiWorkerPool::start() {
     }
 }
 
+void MultiWorkerPool::stop() {
+    for (const auto& pool : _pools) {
+        pool->stop();
+    }
+}
+
 void MultiWorkerPool::submit_publish_version_task(const TAgentTaskRequest& task) {
     auto req = task.publish_version_req;
     for (auto& partition : req.partition_version_infos) {

be/src/agent/multi_worker_pool.h

@@ -10,21 +10,23 @@ namespace starrocks {
 // We create MultiWorkerPool for processing publish version task, these tasks are
 // submitted to one task pool according to its partition id, so the tasks belong to
 // the same partition will be processed by the same worker thread.
-class MultiWorkerPool : public TaskWorkerPool {
+class MultiWorkerPool {
 public:
-    MultiWorkerPool(const TaskWorkerType worker_type, ExecEnv* env, const TMasterInfo& master_info, int worker_num);
+    MultiWorkerPool(const TaskWorkerPool::TaskWorkerType worker_type, ExecEnv* env, const TMasterInfo& master_info,
+                    int worker_num);
 
-    ~MultiWorkerPool() override = default;
-    ;
+    ~MultiWorkerPool() = default;
 
-    void start() override;
+    void start();
+
+    void stop();
 
     // submit task to queue and wait to be executed
-    void submit_task(const TAgentTaskRequest& task) override;
+    void submit_task(const TAgentTaskRequest& task);
 
 private:
     void submit_publish_version_task(const TAgentTaskRequest& task);
 
     std::vector<std::shared_ptr<TaskWorkerPool>> _pools;
 };
-} // namespace starrocks
+} // namespace starrocks