From 5c276fa547c0b5ba9a7a7332968ecc89f2b51105 Mon Sep 17 00:00:00 2001
From: wangjiaming0909 <604227650@qq.com>
Date: Fri, 11 Aug 2023 17:15:17 +0800
Subject: [PATCH] fix: make kill query work for sysscanoperator

---
 source/libs/executor/inc/querytask.h       |  2 +-
 source/libs/executor/inc/tsort.h           |  3 ++-
 source/libs/executor/src/querytask.c       |  2 +-
 source/libs/executor/src/scanoperator.c    |  5 +++--
 source/libs/executor/src/sysscanoperator.c |  5 +++++
 source/libs/executor/src/tsort.c           | 13 ++++++++++---
 6 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/source/libs/executor/inc/querytask.h b/source/libs/executor/inc/querytask.h
index 7241b015a0..0742b9ba4c 100644
--- a/source/libs/executor/inc/querytask.h
+++ b/source/libs/executor/inc/querytask.h
@@ -99,7 +99,7 @@ struct SExecTaskInfo {
 void           buildTaskId(uint64_t taskId, uint64_t queryId, char* dst);
 SExecTaskInfo* doCreateTask(uint64_t queryId, uint64_t taskId, int32_t vgId, EOPTR_EXEC_MODEL model, SStorageAPI* pAPI);
 void           doDestroyTask(SExecTaskInfo* pTaskInfo);
-bool           isTaskKilled(SExecTaskInfo* pTaskInfo);
+bool           isTaskKilled(void* pTaskInfo);
 void           setTaskKilled(SExecTaskInfo* pTaskInfo, int32_t rspCode);
 void           setTaskStatus(SExecTaskInfo* pTaskInfo, int8_t status);
 int32_t        createExecTaskInfo(SSubplan* pPlan, SExecTaskInfo** pTaskInfo, SReadHandle* pHandle, uint64_t taskId,
diff --git a/source/libs/executor/inc/tsort.h b/source/libs/executor/inc/tsort.h
index 57c8bce275..3180173ca7 100644
--- a/source/libs/executor/inc/tsort.h
+++ b/source/libs/executor/inc/tsort.h
@@ -191,7 +191,8 @@ int32_t getProperSortPageSize(size_t rowSize, uint32_t numOfCols);
 bool tsortIsClosed(SSortHandle* pHandle);
 void tsortSetClosed(SSortHandle* pHandle);
 
-void setSingleTableMerge(SSortHandle* pHandle);
+void tsortSetSingleTableMerge(SSortHandle* pHandle);
+void tsortSetAbortCheckFn(SSortHandle* pHandle, bool (*checkFn)(void* param), void* param);
 
 #ifdef __cplusplus
 }
diff --git a/source/libs/executor/src/querytask.c b/source/libs/executor/src/querytask.c
index 22d171e74a..980ef1a61a 100644
--- a/source/libs/executor/src/querytask.c
+++ b/source/libs/executor/src/querytask.c
@@ -59,7 +59,7 @@ SExecTaskInfo* doCreateTask(uint64_t queryId, uint64_t taskId, int32_t vgId, EOP
   return pTaskInfo;
 }
 
-bool isTaskKilled(SExecTaskInfo* pTaskInfo) { return (0 != pTaskInfo->code); }
+bool isTaskKilled(void* pTaskInfo) { return (0 != ((SExecTaskInfo*)pTaskInfo)->code); }
 
 void setTaskKilled(SExecTaskInfo* pTaskInfo, int32_t rspCode) {
   pTaskInfo->code = rspCode;
diff --git a/source/libs/executor/src/scanoperator.c b/source/libs/executor/src/scanoperator.c
index 71b0747be8..a3c5a76a7f 100644
--- a/source/libs/executor/src/scanoperator.c
+++ b/source/libs/executor/src/scanoperator.c
@@ -2928,8 +2928,9 @@ int32_t startGroupTableMergeScan(SOperatorInfo* pOperator) {
     int32_t numOfBufPage = pInfo->sortBufSize / pInfo->bufPageSize;
     pInfo->pSortHandle = tsortCreateSortHandle(pInfo->pSortInfo, SORT_BLOCK_TS_MERGE, pInfo->bufPageSize, numOfBufPage,
                                               pInfo->pSortInputBlock, pTaskInfo->id.str, 0, 0, 0);
-                                          
+
     tsortSetMergeLimit(pInfo->pSortHandle, mergeLimit);
+    tsortSetAbortCheckFn(pInfo->pSortHandle, isTaskKilled, pOperator->pTaskInfo);
   }
 
   tsortSetFetchRawDataFp(pInfo->pSortHandle, getBlockForTableMergeScan, NULL, NULL);
@@ -2949,7 +2950,7 @@ int32_t startGroupTableMergeScan(SOperatorInfo* pOperator) {
 
   int32_t code = TSDB_CODE_SUCCESS;
   if (numOfTable == 1) {
-    setSingleTableMerge(pInfo->pSortHandle);
+    tsortSetSingleTableMerge(pInfo->pSortHandle);
   } else {
     code = tsortOpen(pInfo->pSortHandle);
   }
diff --git a/source/libs/executor/src/sysscanoperator.c b/source/libs/executor/src/sysscanoperator.c
index a1f83dda2f..9048dd43d7 100644
--- a/source/libs/executor/src/sysscanoperator.c
+++ b/source/libs/executor/src/sysscanoperator.c
@@ -1601,6 +1601,11 @@ static SSDataBlock* doSysTableScan(SOperatorInfo* pOperator) {
   SSysTableScanInfo* pInfo = pOperator->info;
   char               dbName[TSDB_DB_NAME_LEN] = {0};
 
+  if (isTaskKilled(pOperator->pTaskInfo)) {
+    setOperatorCompleted(pOperator);
+    return NULL;
+  }
+
   blockDataCleanup(pInfo->pRes);
 
   const char* name = tNameGetTableName(&pInfo->name);
diff --git a/source/libs/executor/src/tsort.c b/source/libs/executor/src/tsort.c
index 1891e93c61..6c4a780dfb 100644
--- a/source/libs/executor/src/tsort.c
+++ b/source/libs/executor/src/tsort.c
@@ -71,12 +71,20 @@ struct SSortHandle {
   SMultiwayMergeTreeInfo* pMergeTree;
 
   bool singleTableMerge;
+
+  bool (*abortCheckFn)(void* param);
+  void* abortCheckParam;
 };
 
-void setSingleTableMerge(SSortHandle* pHandle) {
+void tsortSetSingleTableMerge(SSortHandle* pHandle) {
   pHandle->singleTableMerge = true;
 }
 
+void tsortSetAbortCheckFn(SSortHandle *pHandle, bool (*checkFn)(void *), void* param) {
+  pHandle->abortCheckFn = checkFn;
+  pHandle->abortCheckParam = param;
+}
+
 static int32_t msortComparFn(const void* pLeft, const void* pRight, void* param);
 
 // | offset[0] | offset[1] |....| nullbitmap | data |...|
@@ -726,11 +734,10 @@ static int32_t doInternalMergeSort(SSortHandle* pHandle) {
 
       SArray* pPageIdList = taosArrayInit(4, sizeof(int32_t));
       while (1) {
-        if (tsortIsClosed(pHandle)) {
+        if (tsortIsClosed(pHandle) || (pHandle->abortCheckFn && pHandle->abortCheckFn(pHandle->abortCheckParam))) {
           code = terrno = TSDB_CODE_TSC_QUERY_CANCELLED;
           return code;
         }
-
         SSDataBlock* pDataBlock = getSortedBlockDataInner(pHandle, &pHandle->cmpParam, numOfRows);
         if (pDataBlock == NULL) {
           break;