diff --git a/be/src/service/internal_service.cpp b/be/src/service/internal_service.cpp index f68716eebc863c..f118e8b412211c 100644 --- a/be/src/service/internal_service.cpp +++ b/be/src/service/internal_service.cpp @@ -671,8 +671,7 @@ void PInternalService::fetch_arrow_data(google::protobuf::RpcController* control PFetchArrowDataResult* result, google::protobuf::Closure* done) { bool ret = _arrow_flight_work_pool.try_offer([request, result, done]() { - brpc::ClosureGuard closure_guard(done); - auto ctx = vectorized::GetArrowResultBatchCtx::create_shared(result); + auto ctx = vectorized::GetArrowResultBatchCtx::create_shared(result, done); TUniqueId unique_id = UniqueId(request->finst_id()).to_thrift(); // query_id or instance_id std::shared_ptr arrow_buffer; auto st = ExecEnv::GetInstance()->result_mgr()->find_buffer(unique_id, arrow_buffer); diff --git a/be/src/vec/sink/varrow_flight_result_writer.cpp b/be/src/vec/sink/varrow_flight_result_writer.cpp index b828078554624c..6c9c393ae4a8fa 100644 --- a/be/src/vec/sink/varrow_flight_result_writer.cpp +++ b/be/src/vec/sink/varrow_flight_result_writer.cpp @@ -33,6 +33,7 @@ namespace doris::vectorized { void GetArrowResultBatchCtx::on_failure(const Status& status) { DCHECK(!status.ok()) << "status is ok, errmsg=" << status; status.to_protobuf(_result->mutable_status()); + _done->Run(); } void GetArrowResultBatchCtx::on_close(int64_t packet_seq, int64_t /* returned_rows */) { @@ -40,6 +41,7 @@ void GetArrowResultBatchCtx::on_close(int64_t packet_seq, int64_t /* returned_ro status.to_protobuf(_result->mutable_status()); _result->set_packet_seq(packet_seq); _result->set_eos(true); + _done->Run(); } Status GetArrowResultBatchCtx::on_data(const std::shared_ptr& block, @@ -72,6 +74,8 @@ Status GetArrowResultBatchCtx::on_data(const std::shared_ptr& _result->clear_block(); } st.to_protobuf(_result->mutable_status()); + + _done->Run(); return Status::OK(); } diff --git a/be/src/vec/sink/varrow_flight_result_writer.h b/be/src/vec/sink/varrow_flight_result_writer.h index f84c3a12fa8dbe..d1f5e21404a07f 100644 --- a/be/src/vec/sink/varrow_flight_result_writer.h +++ b/be/src/vec/sink/varrow_flight_result_writer.h @@ -35,7 +35,8 @@ class GetArrowResultBatchCtx { public: using ResultType = vectorized::Block; ENABLE_FACTORY_CREATOR(GetArrowResultBatchCtx) - GetArrowResultBatchCtx(PFetchArrowDataResult* result) : _result(result) {} + GetArrowResultBatchCtx(PFetchArrowDataResult* result, google::protobuf::Closure* done) + : _result(result), _done(done) {} #ifdef BE_TEST GetArrowResultBatchCtx() = default; #endif @@ -53,6 +54,7 @@ class GetArrowResultBatchCtx { int32_t _max_msg_size = std::numeric_limits::max(); #endif PFetchArrowDataResult* _result = nullptr; + google::protobuf::Closure* _done = nullptr; }; class ArrowFlightResultBlockBuffer final : public ResultBlockBuffer { diff --git a/be/test/vec/sink/arrow_result_block_buffer_test.cpp b/be/test/vec/sink/arrow_result_block_buffer_test.cpp index 48718361185eb4..7fb1d8adf1f7d2 100644 --- a/be/test/vec/sink/arrow_result_block_buffer_test.cpp +++ b/be/test/vec/sink/arrow_result_block_buffer_test.cpp @@ -34,6 +34,14 @@ namespace doris::vectorized { +class MockClosure : public google::protobuf::Closure { +public: + MockClosure() {} + MockClosure(std::function cb) : _cb(cb) {} + void Run() override { _cb(); } + + std::function _cb; +}; class ArrowResultBlockBufferTest : public ::testing::Test { public: ArrowResultBlockBufferTest() = default; @@ -44,8 +52,9 @@ class MockGetArrowResultBatchCtx : public GetArrowResultBatchCtx { public: ENABLE_FACTORY_CREATOR(MockGetArrowResultBatchCtx) MockGetArrowResultBatchCtx(std::function fail_cb, std::function close_cb, - std::function data_cb, PFetchArrowDataResult* result) - : GetArrowResultBatchCtx(result), + std::function data_cb, PFetchArrowDataResult* result, + google::protobuf::Closure* done) + : GetArrowResultBatchCtx(result, done), _fail_cb(fail_cb), _close_cb(close_cb), _data_cb(data_cb) {} @@ -78,9 +87,11 @@ TEST_F(ArrowResultBlockBufferTest, TestArrowResultBlockBuffer) { ArrowFlightResultBlockBuffer buffer(TUniqueId(), &state, schema, buffer_size); buffer.set_dependency(ins_id, dep); PFetchArrowDataResult presult; + + MockClosure done([&]() -> void { std::cout << "cb" << std::endl; }); std::shared_ptr ctx = MockGetArrowResultBatchCtx::create_shared( [&]() -> void { fail = true; }, [&]() -> void { close = true; }, - [&]() -> void { data = true; }, &presult); + [&]() -> void { data = true; }, &presult, &done); { auto num_rows = 2; @@ -201,9 +212,11 @@ TEST_F(ArrowResultBlockBufferTest, TestCancelArrowResultBlockBuffer) { ArrowFlightResultBlockBuffer buffer(TUniqueId(), &state, schema, buffer_size); buffer.set_dependency(ins_id, dep); PFetchArrowDataResult presult; + + MockClosure done([&]() -> void { std::cout << "cb" << std::endl; }); std::shared_ptr ctx = MockGetArrowResultBatchCtx::create_shared( [&]() -> void { fail = true; }, [&]() -> void { close = true; }, - [&]() -> void { data = true; }, &presult); + [&]() -> void { data = true; }, &presult, &done); { EXPECT_TRUE(buffer.get_batch(ctx).ok()); @@ -273,9 +286,11 @@ TEST_F(ArrowResultBlockBufferTest, TestErrorClose) { ArrowFlightResultBlockBuffer buffer(TUniqueId(), &state, schema, buffer_size); buffer.set_dependency(ins_id, dep); PFetchArrowDataResult presult; + + MockClosure done([&]() -> void { std::cout << "cb" << std::endl; }); std::shared_ptr ctx = MockGetArrowResultBatchCtx::create_shared( [&]() -> void { fail = true; }, [&]() -> void { close = true; }, - [&]() -> void { data = true; }, &presult); + [&]() -> void { data = true; }, &presult, &done); { EXPECT_TRUE(buffer.get_batch(ctx).ok()); @@ -330,9 +345,11 @@ TEST_F(ArrowResultBlockBufferTest, TestArrowResultSerializeFailure) { ArrowFlightResultBlockBuffer buffer(TUniqueId(), &state, schema, buffer_size); buffer.set_dependency(ins_id, dep); PFetchArrowDataResult presult; + + MockClosure done([&]() -> void { std::cout << "cb" << std::endl; }); std::shared_ptr ctx = MockGetArrowResultBatchCtx::create_shared( [&]() -> void { fail = true; }, [&]() -> void { close = true; }, - [&]() -> void { data = true; }, &presult); + [&]() -> void { data = true; }, &presult, &done); { auto num_rows = 2; diff --git a/be/test/vec/sink/get_result_batch_test.cpp b/be/test/vec/sink/get_result_batch_test.cpp index 82eff9dd289806..b820bd43ecf73d 100644 --- a/be/test/vec/sink/get_result_batch_test.cpp +++ b/be/test/vec/sink/get_result_batch_test.cpp @@ -40,10 +40,10 @@ class GetResultBatchCtxTest : public ::testing::Test { class MockClosure : public google::protobuf::Closure { public: + MockClosure() {} MockClosure(std::function cb) : _cb(cb) {} void Run() override { _cb(); } -private: std::function _cb; }; @@ -126,7 +126,9 @@ TEST_F(GetResultBatchCtxTest, TestGetResultBatchCtx) { TEST_F(GetResultBatchCtxTest, TestGetArrowResultBatchCtx) { PFetchArrowDataResult result; - auto ctx = GetArrowResultBatchCtx::create_shared(&result); + MockClosure closure; + closure._cb = [&]() { std::cout << "cb" << std::endl; }; + auto ctx = GetArrowResultBatchCtx::create_shared(&result, &closure); { // on_failure