// 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 "exec/exchange/vdata_stream_mgr.h"

#include <gen_cpp/Types_types.h>

#include <gen_cpp/data.pb.h>

#include <gen_cpp/internal_service.pb.h>

#include <gen_cpp/types.pb.h>

#include <stddef.h>

#include <memory>

#include <ostream>

#include <string>

#include <vector>

#include "common/logging.h"

#include "exec/exchange/vdata_stream_recvr.h"

#include "util/hash_util.hpp"

namespace doris {

#include "common/compile_check_begin.h"

VDataStreamMgr::VDataStreamMgr() {

    // TODO: metric

}

VDataStreamMgr::~VDataStreamMgr() {

    // Has to call close here, because receiver will check if the receiver is closed.

    // It will core during graceful stop.

    auto receivers = std::vector<std::shared_ptr<VDataStreamRecvr>>();

    {

        std::shared_lock l(_lock);

        auto receiver_iterator = _receiver_map.begin();

        while (receiver_iterator != _receiver_map.end()) {

            // Could not call close directly, because during close method, it will remove itself

            // from the map, and modify the map, it will core.

            receivers.push_back(receiver_iterator->second);

            receiver_iterator++;

        }

    }

    for (auto iter = receivers.begin(); iter != receivers.end(); ++iter) {

        (*iter)->close();

    }

}

inline uint32_t VDataStreamMgr::get_hash_value(const TUniqueId& fragment_instance_id,

                                               PlanNodeId node_id) {

    uint32_t value = HashUtil::hash(&fragment_instance_id.lo, 8, 0);

    value = HashUtil::hash(&fragment_instance_id.hi, 8, value);

    value = HashUtil::hash(&node_id, 4, value);

    return value;

}

std::shared_ptr<VDataStreamRecvr> VDataStreamMgr::create_recvr(

        RuntimeState* state, RuntimeProfile::HighWaterMarkCounter* memory_used_counter,

        const TUniqueId& fragment_instance_id, PlanNodeId dest_node_id, int num_senders,

        RuntimeProfile* profile, bool is_merging, size_t data_queue_capacity) {

    DCHECK(profile != nullptr);

    VLOG_FILE << "creating receiver for fragment=" << print_id(fragment_instance_id)

              << ", node=" << dest_node_id;

    std::shared_ptr<VDataStreamRecvr> recvr(new VDataStreamRecvr(

            this, memory_used_counter, state, fragment_instance_id, dest_node_id, num_senders,

            is_merging, profile, data_queue_capacity));

    uint32_t hash_value = get_hash_value(fragment_instance_id, dest_node_id);

    std::unique_lock l(_lock);

    _fragment_stream_set.insert(std::make_pair(fragment_instance_id, dest_node_id));

    _receiver_map.insert(std::make_pair(hash_value, recvr));

    return recvr;

}

Status VDataStreamMgr::find_recvr(const TUniqueId& fragment_instance_id, PlanNodeId node_id,

                                  std::shared_ptr<VDataStreamRecvr>* res, bool acquire_lock) {

    VLOG_ROW << "looking up fragment_instance_id=" << print_id(fragment_instance_id)

             << ", node=" << node_id;

    uint32_t hash_value = get_hash_value(fragment_instance_id, node_id);

    // Create lock guard and not own lock currently and will lock conditionally

    std::shared_lock recvr_lock(_lock, std::defer_lock);

    if (acquire_lock) {

        recvr_lock.lock();

    }

    std::pair<StreamMap::iterator, StreamMap::iterator> range =

            _receiver_map.equal_range(hash_value);

    while (range.first != range.second) {

        auto recvr = range.first->second;

        if (recvr->fragment_instance_id() == fragment_instance_id &&

            recvr->dest_node_id() == node_id) {

            *res = recvr;

            return Status::OK();

        }

        ++range.first;

    }

    return Status::InvalidArgument("Could not find local receiver for node {} with instance {}",

                                   node_id, print_id(fragment_instance_id));

}

Status VDataStreamMgr::transmit_block(const PTransmitDataParams* request,

                                      ::google::protobuf::Closure** done,

                                      const int64_t wait_for_worker) {

    const PUniqueId& finst_id = request->finst_id();

    TUniqueId t_finst_id;

    t_finst_id.hi = finst_id.hi();

    t_finst_id.lo = finst_id.lo();

    std::shared_ptr<VDataStreamRecvr> recvr = nullptr;

    ThreadCpuStopWatch cpu_time_stop_watch;

    cpu_time_stop_watch.start();

    static_cast<void>(find_recvr(t_finst_id, request->node_id(), &recvr));

    if (recvr == nullptr) {

        // The receiver may remove itself from the receiver map via deregister_recvr()

        // at any time without considering the remaining number of senders.

        // As a consequence, find_recvr() may return an innocuous NULL if a thread

        // calling deregister_recvr() beat the thread calling find_recvr()

        // in acquiring _lock.

        //

        // e.g. for broadcast join build side, only one instance will build the hash table,

        // all other instances don't need build side data and will close the data stream receiver.

        //

        // TODO: Rethink the lifecycle of DataStreamRecvr to distinguish

        // errors from receiver-initiated teardowns.

        return Status::EndOfFile("data stream receiver closed");

    }

    // Lock the fragment context to ensure the runtime state and other objects are not

    // deconstructed

    auto ctx_lock = recvr->task_exec_ctx();

    if (ctx_lock == nullptr) {

        // Do not return internal error, because when query finished, the downstream node

        // may finish before upstream node. And the object maybe deconstructed. If return error

        // then the upstream node may report error status to FE, the query is failed.

        return Status::EndOfFile("data stream receiver is deconstructed");

    }

    bool eos = request->eos();

    Status exec_status =

            request->has_exec_status() ? Status::create(request->exec_status()) : Status::OK();

    auto sender_id = request->sender_id();

    auto be_number = request->be_number();

    if (!request->blocks().empty()) {

        RETURN_IF_ERROR(recvr->add_blocks(request, done, wait_for_worker,

                                          cpu_time_stop_watch.elapsed_time()));

    } else if (request->has_block()) {

        // old logic, for compatibility

        std::unique_ptr<PBlock> pblock_ptr = std::make_unique<PBlock>();

        pblock_ptr->CopyFrom(request->block());

        RETURN_IF_ERROR(recvr->add_block(std::move(pblock_ptr), sender_id, be_number,

                                         request->packet_seq(), eos ? nullptr : done,

                                         wait_for_worker, cpu_time_stop_watch.elapsed_time()));

    }

    if (eos) {

        recvr->remove_sender(sender_id, be_number, exec_status);

    }

    return Status::OK();

}

Status VDataStreamMgr::deregister_recvr(const TUniqueId& fragment_instance_id, PlanNodeId node_id) {

    std::shared_ptr<VDataStreamRecvr> targert_recvr;

    VLOG_QUERY << "deregister_recvr(): fragment_instance_id=" << print_id(fragment_instance_id)

               << ", node=" << node_id;

    uint32_t hash_value = get_hash_value(fragment_instance_id, node_id);

    {

        std::unique_lock l(_lock);

        auto range = _receiver_map.equal_range(hash_value);

        while (range.first != range.second) {

            const std::shared_ptr<VDataStreamRecvr>& recvr = range.first->second;

            if (recvr->fragment_instance_id() == fragment_instance_id &&

                recvr->dest_node_id() == node_id) {

                targert_recvr = recvr;

                _fragment_stream_set.erase(

                        std::make_pair(recvr->fragment_instance_id(), recvr->dest_node_id()));

                _receiver_map.erase(range.first);

                break;

            }

            ++range.first;

        }

    }

    // Notify concurrent add_data() requests that the stream has been terminated.

    // cancel_stream maybe take a long time, so we handle it out of lock.

    if (targert_recvr) {

        targert_recvr->cancel_stream(Status::OK());

        return Status::OK();

    } else {

        return Status::InternalError("unknown row receiver id: fragment_instance_id={}, node_id={}",

                                     print_id(fragment_instance_id), node_id);

    }

}

void VDataStreamMgr::cancel(const TUniqueId& fragment_instance_id, Status exec_status) {

    VLOG_QUERY << "cancelling all streams for fragment=" << print_id(fragment_instance_id);

    std::vector<std::shared_ptr<VDataStreamRecvr>> recvrs;

    {

        std::shared_lock l(_lock);

        FragmentStreamSet::iterator i =

                _fragment_stream_set.lower_bound(std::make_pair(fragment_instance_id, 0));

        while (i != _fragment_stream_set.end() && i->first == fragment_instance_id) {

            std::shared_ptr<VDataStreamRecvr> recvr;

            WARN_IF_ERROR(find_recvr(i->first, i->second, &recvr, false), "");

            if (recvr == nullptr) {

                // keep going but at least log it

                std::stringstream err;

                err << "cancel(): missing in stream_map: fragment=" << print_id(i->first)

                    << " node=" << i->second;

                LOG(ERROR) << err.str();

            } else {

                recvrs.push_back(recvr);

            }

            ++i;

        }

    }

    // cancel_stream maybe take a long time, so we handle it out of lock.

    for (auto& it : recvrs) {

        it->cancel_stream(exec_status);

    }

}

} // namespace doris