update

apache · Dec 31, 2023 · a8a53b5 · a8a53b5
1 parent ecb3486
commit a8a53b5
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 6 deletions.
diff --git a/be/src/vec/sink/load_stream_stub.h b/be/src/vec/sink/load_stream_stub.h
@@ -110,7 +110,10 @@ class LoadStreamStub {
                     return Status::OK();
                 }
                 if (_stub->is_cancelled()) {
-                    return Status::Cancelled(_stub->cancel_reason());
+                    return _stub->cancelled_state();
+                }
+                if (_stub->runtime_state()->is_cancelled()) {
+                    return Status::Cancelled(_stub->runtime_state()->cancel_reason());
                 }
                 // wait 1s once time.
                 ret = _close_cv.wait_for(lock, 1);
@@ -159,13 +162,17 @@ class LoadStreamStub {
     // copy constructor, shared_ptr members are shared
     LoadStreamStub(LoadStreamStub& stub, RuntimeState* state);
 
-    bool is_cancelled() const {
-        if (_state == nullptr) {
-            return false;
-        }
-        return _state->is_cancelled();
+    void cancel(Status status) {
+        _cancel = true;
+        _cancel_status = status;
     }
 
+    RuntimeState* runtime_state() const { return _state; }
+
+    bool is_cancelled() const { return _cancel; }
+
+    Status cancelled_state() const { return _cancel_status; }
+
     std::string cancel_reason() const {
         if (_state == nullptr) {
             return "";
@@ -263,6 +270,9 @@ class LoadStreamStub {
 protected:
     RuntimeState* _state = nullptr;
 
+    bool _cancel = false;
+    Status _cancel_status;
+
     std::atomic<bool> _is_init;
     bthread::Mutex _mutex;
 

diff --git a/be/src/vec/sink/load_stream_stub_pool.cpp b/be/src/vec/sink/load_stream_stub_pool.cpp
@@ -54,6 +54,12 @@ void LoadStreams::release() {
     }
 }
 
+void LoadStreams::cancel(Status status) {
+    for (auto& stream : _streams) {
+        stream->cancel(status);
+    }
+}
+
 LoadStreamStubPool::LoadStreamStubPool() = default;
 
 LoadStreamStubPool::~LoadStreamStubPool() = default;

diff --git a/be/src/vec/sink/load_stream_stub_pool.h b/be/src/vec/sink/load_stream_stub_pool.h
@@ -78,6 +78,8 @@ class LoadStreams {
 
     void release();
 
+    void cancel(Status status);
+
     Streams& streams() { return _streams; }
 
 private:

diff --git a/be/src/vec/sink/writer/vtablet_writer_v2.cpp b/be/src/vec/sink/writer/vtablet_writer_v2.cpp
@@ -457,6 +457,7 @@ Status VTabletWriterV2::_cancel(Status status) {
     }
     for (const auto& [_, streams] : _streams_for_node) {
         streams->release();
+        streams->cancel(status);
     }
     return Status::OK();
 }