// 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.

#include "load/routine_load/data_consumer_pool.h"

#include <gen_cpp/Types_types.h>

#include <algorithm>

// IWYU pragma: no_include <bits/chrono.h>

#include <chrono> // IWYU pragma: keep

#include <ctime>

#include <iterator>

#include <map>

#include <ostream>

#include <vector>

#include "common/config.h"

#include "common/logging.h"

#include "common/status.h"

#include "load/routine_load/data_consumer.h"

#include "load/routine_load/data_consumer_group.h"

#include "load/stream_load/stream_load_context.h"

#include "util/uid_util.h"

namespace doris {

Status DataConsumerPool::get_consumer(std::shared_ptr<StreamLoadContext> ctx,

std::shared_ptr<DataConsumer>* ret) {

std::unique_lock<std::mutex> l(_lock);

// check if there is an available consumer.

// if has, return it, also remove it from the pool

auto iter = std::begin(_pool);

while (iter != std::end(_pool)) {

if ((*iter)->match(ctx)) {

VLOG_NOTICE << "get an available data consumer from pool: " << (*iter)->id();

static_cast<void>((*iter)->reset());

*ret = *iter;

iter = _pool.erase(iter);

return Status::OK();

} else {

++iter;

}

}

// no available consumer, create a new one

std::shared_ptr<DataConsumer> consumer;

switch (ctx->load_src_type) {

case TLoadSourceType::KAFKA:

consumer = std::make_shared<KafkaDataConsumer>(ctx);

break;

case TLoadSourceType::KINESIS:

consumer = std::make_shared<KinesisDataConsumer>(ctx);

break;

default:

return Status::InternalError("PAUSE: unknown routine load task type: {}", ctx->load_type);

}

// init the consumer

RETURN_IF_ERROR(consumer->init(ctx));

VLOG_NOTICE << "create new data consumer: " << consumer->id();

*ret = consumer;

return Status::OK();

}

Status DataConsumerPool::get_consumer_grp(std::shared_ptr<StreamLoadContext> ctx,

std::shared_ptr<DataConsumerGroup>* ret) {

size_t consumer_num = config::max_consumer_num_per_group;

std::shared_ptr<DataConsumerGroup> grp;

switch (ctx->load_src_type) {

case TLoadSourceType::KAFKA: {

DCHECK(ctx->kafka_info);

if (ctx->kafka_info->begin_offset.size() == 0) {

return Status::InternalError(

"PAUSE: The size of begin_offset of task should not be 0.");

}

consumer_num = std::min(consumer_num, ctx->kafka_info->begin_offset.size());

grp = std::make_shared<KafkaDataConsumerGroup>(consumer_num);

break;

}

case TLoadSourceType::KINESIS: {

DCHECK(ctx->kinesis_info);

if (ctx->kinesis_info->begin_sequence_number.size() == 0) {

return Status::InternalError(

"PAUSE: The size of begin_sequence_number of task should not be 0.");

}

consumer_num = std::min(consumer_num, ctx->kinesis_info->begin_sequence_number.size());

grp = std::make_shared<KinesisDataConsumerGroup>(consumer_num);

break;

}

default:

return Status::Cancelled("unknown routine load task type: {}", ctx->load_type);

}

for (int i = 0; i < consumer_num; ++i) {

std::shared_ptr<DataConsumer> consumer;

RETURN_IF_ERROR(get_consumer(ctx, &consumer));

grp->add_consumer(consumer);

}

LOG(INFO) << "get consumer group " << grp->grp_id() << " with " << consumer_num << " consumers";

*ret = grp;

return Status::OK();

}

void DataConsumerPool::return_consumer(std::shared_ptr<DataConsumer> consumer) {

std::unique_lock<std::mutex> l(_lock);

if (_pool.size() == config::routine_load_consumer_pool_size) {

VLOG_NOTICE << "data consumer pool is full: " << _pool.size() << "-"

<< config::routine_load_consumer_pool_size

<< ", discard the returned consumer: " << consumer->id();

return;

}

static_cast<void>(consumer->reset());

_pool.push_back(consumer);

VLOG_NOTICE << "return the data consumer: " << consumer->id()

<< ", current pool size: " << _pool.size();

}

void DataConsumerPool::return_consumers(DataConsumerGroup* grp) {

for (std::shared_ptr<DataConsumer> consumer : grp->consumers()) {

return_consumer(consumer);

}

}

Status DataConsumerPool::start_bg_worker() {

RETURN_IF_ERROR(Thread::create(

"ResultBufferMgr", "clean_idle_consumer",

[this]() {

do {

_clean_idle_consumer_bg();

} while (!_stop_background_threads_latch.wait_for(std::chrono::seconds(60)));

},

&_clean_idle_consumer_thread));

return Status::OK();

}

void DataConsumerPool::_clean_idle_consumer_bg() {

const static int32_t max_idle_time_second = 600;

std::unique_lock<std::mutex> l(_lock);

time_t now = time(nullptr);

auto iter = std::begin(_pool);

while (iter != std::end(_pool)) {

if (difftime(now, (*iter)->last_visit_time()) >= max_idle_time_second) {

LOG(INFO) << "remove data consumer " << (*iter)->id()

<< ", since it last visit: " << (*iter)->last_visit_time()

<< ", now: " << now;

iter = _pool.erase(iter);

} else {

++iter;

}

}

}

} // end namespace doris