Files
doris/be/src/runtime/buffer_control_block.h
Zhengguo Yang 146060dfc0 [Bug]Fix result_writer may coredump (#6482)
fix result_writer may coredump, let BufferControlBlock owns the memory
2021-08-22 22:04:00 +08:00

132 lines
4.4 KiB
C++

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#ifndef DORIS_BE_RUNTIME_BUFFER_CONTROL_BLOCK_H
#define DORIS_BE_RUNTIME_BUFFER_CONTROL_BLOCK_H
#include <condition_variable>
#include <deque>
#include <list>
#include <mutex>
#include "common/status.h"
#include "gen_cpp/Types_types.h"
#include "runtime/query_statistics.h"
namespace google {
namespace protobuf {
class Closure;
}
} // namespace google
namespace brpc {
class Controller;
}
namespace doris {
class TFetchDataResult;
class PFetchDataResult;
struct GetResultBatchCtx {
brpc::Controller* cntl = nullptr;
// In version 0.15, we change brpc client from jprotobuf to grpc.
// And the response data is moved from rpc attachment to protobuf boby.
// This variables is for backwards compatibility when upgrading Doris.
// If set to true, the response data is still transferred as attachment.
// If set to false, the response data is transferred in protobuf body.
bool resp_in_attachment = true;
PFetchDataResult* result = nullptr;
google::protobuf::Closure* done = nullptr;
GetResultBatchCtx(brpc::Controller* cntl_, bool resp_in_attachment_, PFetchDataResult* result_,
google::protobuf::Closure* done_)
: cntl(cntl_), resp_in_attachment(resp_in_attachment_), result(result_), done(done_) {}
void on_failure(const Status& status);
void on_close(int64_t packet_seq, QueryStatistics* statistics = nullptr);
void on_data(const std::unique_ptr<TFetchDataResult>& t_result, int64_t packet_seq, bool eos = false);
};
// buffer used for result customer and producer
class BufferControlBlock {
public:
BufferControlBlock(const TUniqueId& id, int buffer_size);
~BufferControlBlock();
Status init();
Status add_batch(std::unique_ptr<TFetchDataResult>& result);
// get result from batch, use timeout?
Status get_batch(TFetchDataResult* result);
void get_batch(GetResultBatchCtx* ctx);
// close buffer block, set _status to exec_status and set _is_close to true;
// called because data has been read or error happened.
Status close(Status exec_status);
// this is called by RPC, called from coordinator
Status cancel();
const TUniqueId& fragment_id() const { return _fragment_id; }
void set_query_statistics(std::shared_ptr<QueryStatistics> statistics) {
_query_statistics = statistics;
}
void update_num_written_rows(int64_t num_rows) {
// _query_statistics may be null when the result sink init failed
// or some other failure.
// and the number of written rows is only needed when all things go well.
if (_query_statistics.get() != nullptr) {
_query_statistics->set_returned_rows(num_rows);
}
}
private:
typedef std::list<std::unique_ptr<TFetchDataResult>> ResultQueue;
// result's query id
TUniqueId _fragment_id;
bool _is_close;
bool _is_cancelled;
Status _status;
int _buffer_rows;
int _buffer_limit;
int64_t _packet_num;
// blocking queue for batch
ResultQueue _batch_queue;
// protects all subsequent data in this block
std::mutex _lock;
// signal arrival of new batch or the eos/cancelled condition
std::condition_variable _data_arrival;
// signal removal of data by stream consumer
std::condition_variable _data_removal;
std::deque<GetResultBatchCtx*> _waiting_rpc;
// It is shared with PlanFragmentExecutor and will be called in two different
// threads. But their calls are all at different time, there is no problem of
// multithreading access.
std::shared_ptr<QueryStatistics> _query_statistics;
};
} // namespace doris
#endif