fix a bug that ExchangeReceiver can't be canceled (#4441) (#4445)

close #4432

dbms/src/DataStreams/TiRemoteBlockInputStream.h

@@ -241,6 +241,7 @@ class TiRemoteBlockInputStream : public IProfilingBlockInputStream
     virtual void readSuffix() override
     {
         LOG_FMT_DEBUG(log, "finish read {} rows from remote", total_rows);
+        remote_reader->close();
     }
 };
 

dbms/src/Flash/Coprocessor/CoprocessorReader.h

@@ -95,6 +95,7 @@ class CoprocessorReader
 
     void cancel() { resp_iter.cancel(); }
 
+
     static DecodeDetail decodeChunks(
         const std::shared_ptr<tipb::SelectResponse> & resp,
         std::queue<Block> & block_queue,
@@ -187,6 +188,8 @@ class CoprocessorReader
         collected = false;
     }
 
+    void close() {}
+
     bool collected = false;
     int concurrency_;
 };

dbms/src/Flash/Coprocessor/DAGContext.cpp

@@ -219,6 +219,14 @@ const std::unordered_map<String, std::shared_ptr<ExchangeReceiver>> & DAGContext
     return mpp_exchange_receiver_map;
 }
 
+void DAGContext::cancelAllExchangeReceiver()
+{
+    for (auto & it : mpp_exchange_receiver_map)
+    {
+        it.second->cancel();
+    }
+}
+
 int DAGContext::getNewThreadCountOfExchangeReceiver() const
 {
     return new_thread_count_of_exchange_receiver;

dbms/src/Flash/Coprocessor/DAGContext.h

@@ -266,6 +266,8 @@ class DAGContext
         return sql_mode & f;
     }
 
+    void cancelAllExchangeReceiver();
+
     void initExchangeReceiverIfMPP(Context & context, size_t max_streams);
     const std::unordered_map<String, std::shared_ptr<ExchangeReceiver>> & getMPPExchangeReceiverMap() const;
 

dbms/src/Flash/Mpp/ExchangeReceiver.cpp

@@ -310,15 +310,21 @@ ExchangeReceiverBase<RPCContext>::ExchangeReceiverBase(
 template <typename RPCContext>
 ExchangeReceiverBase<RPCContext>::~ExchangeReceiverBase()
 {
-    setState(ExchangeReceiverState::CLOSED);
-    msg_channel.finish();
+    close();
     thread_manager->wait();
 }
 
 template <typename RPCContext>
 void ExchangeReceiverBase<RPCContext>::cancel()
 {
-    setState(ExchangeReceiverState::CANCELED);
+    setEndState(ExchangeReceiverState::CANCELED);
+    msg_channel.finish();
+}
+
+template <typename RPCContext>
+void ExchangeReceiverBase<RPCContext>::close()
+{
+    setEndState(ExchangeReceiverState::CLOSED);
     msg_channel.finish();
 }
 
@@ -599,10 +605,16 @@ ExchangeReceiverResult ExchangeReceiverBase<RPCContext>::nextResult(std::queue<B
 }
 
 template <typename RPCContext>
-void ExchangeReceiverBase<RPCContext>::setState(ExchangeReceiverState new_state)
+bool ExchangeReceiverBase<RPCContext>::setEndState(ExchangeReceiverState new_state)
 {
+    assert(new_state == ExchangeReceiverState::CANCELED || new_state == ExchangeReceiverState::CLOSED);
     std::unique_lock lock(mu);
+    if (state == ExchangeReceiverState::CANCELED || state == ExchangeReceiverState::CLOSED)
+    {
+        return false;
+    }
     state = new_state;
+    return true;
 }
 
 template <typename RPCContext>

dbms/src/Flash/Mpp/ExchangeReceiver.h

@@ -98,6 +98,8 @@ class ExchangeReceiverBase
 
     void cancel();
 
+    void close();
+
     const DAGSchema & getOutputSchema() const { return schema; }
 
     ExchangeReceiverResult nextResult(
@@ -129,7 +131,7 @@ class ExchangeReceiverBase
     void readLoop(const Request & req);
     void reactor(const std::vector<Request> & async_requests);
 
-    void setState(ExchangeReceiverState new_state);
+    bool setEndState(ExchangeReceiverState new_state);
     ExchangeReceiverState getState();
 
     DecodeDetail decodeChunks(

dbms/src/Flash/Mpp/MPPTask.cpp

@@ -353,6 +353,8 @@ void MPPTask::runImpl()
     else
     {
         context->getProcessList().sendCancelToQuery(context->getCurrentQueryId(), context->getClientInfo().current_user, true);
+        if (dag_context)
+            dag_context->cancelAllExchangeReceiver();
         writeErrToAllTunnels(err_msg);
     }
     LOG_FMT_INFO(log, "task ends, time cost is {} ms.", stopwatch.elapsedMilliseconds());