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

#pragma once

#include <gen_cpp/Types_types.h>

#include <glog/logging.h>

#include <google/protobuf/stubs/callback.h>

#include <atomic>

#include <condition_variable>

#include <cstddef>

#include <cstdint>

#include <deque>

#include <list>

#include <memory>

#include <mutex>

#include <ostream>

#include <sstream>

#include <thread>

#include <unordered_map>

#include <unordered_set>

#include <utility>

#include <vector>

#include "common/config.h"

#include "common/global_types.h"

#include "common/object_pool.h"

#include "common/status.h"

#include "core/block/block.h"

#include "exprs/vexpr_fwd.h"

#include "runtime/descriptors.h"

#include "runtime/runtime_profile.h"

#include "runtime/task_execution_context.h"

#include "runtime/thread_context.h"

#include "util/stopwatch.hpp"

namespace doris {

class MemTracker;

class PBlock;

class PTransmitDataParams;

class MemTrackerLimiter;

class RuntimeState;

class Dependency;

class ExchangeLocalState;

class VDataStreamMgr;

class VSortedRunMerger;

class VDataStreamRecvr;

class VDataStreamRecvr : public HasTaskExecutionCtx {

public:

    class SenderQueue;

    VDataStreamRecvr(VDataStreamMgr* stream_mgr, RuntimeProfile::HighWaterMarkCounter* counter,

                     RuntimeState* state, const TUniqueId& fragment_instance_id,

                     PlanNodeId dest_node_id, int num_senders, bool is_merging,

                     RuntimeProfile* profile, size_t data_queue_capacity);

    ~VDataStreamRecvr() override;

    MOCK_FUNCTION Status create_merger(const VExprContextSPtrs& ordering_expr,

                                       const std::vector<bool>& is_asc_order,

                                       const std::vector<bool>& nulls_first, size_t batch_size,

                                       int64_t limit, size_t offset);

    std::vector<SenderQueue*> sender_queues() const { return _sender_queues; }

    Status add_block(std::unique_ptr<PBlock> pblock, int sender_id, int be_number,

                     int64_t packet_seq, ::google::protobuf::Closure** done,

                     const int64_t wait_for_worker, const uint64_t time_to_find_recvr);

    Status add_blocks(const PTransmitDataParams* request, ::google::protobuf::Closure** done,

                      const int64_t wait_for_worker, const uint64_t time_to_find_recvr);

    void add_block(Block* block, int sender_id, bool use_move);

    std::string debug_string();

    MOCK_FUNCTION Status get_next(Block* block, bool* eos);

    const TUniqueId& fragment_instance_id() const { return _fragment_instance_id; }

    PlanNodeId dest_node_id() const { return _dest_node_id; }

    // Indicate that a particular sender is done. Delegated to the appropriate

    // sender queue. Called from DataStreamMgr.

    void remove_sender(int sender_id, int be_number, Status exec_status);

    void cancel_stream(Status exec_status);

    MOCK_FUNCTION void close();

    // When the source reaches eos = true

    void set_sink_dep_always_ready() const;

    // Careful: stream sender will call this function for a local receiver,

    // accessing members of receiver that are allocated by Object pool

    // in this function is not safe.

    MOCK_FUNCTION bool exceeds_limit(size_t block_byte_size);

    bool queue_exceeds_limit(size_t byte_size) const;

    bool is_closed() const { return _is_closed; }

    std::shared_ptr<Dependency> get_local_channel_dependency(int sender_id);

    void set_low_memory_mode() { _sender_queue_mem_limit = 1012 * 1024; }

private:

    friend struct BlockSupplierSortCursorImpl;

    // DataStreamMgr instance used to create this recvr. (Not owned)

    VDataStreamMgr* _mgr = nullptr;

    RuntimeProfile::HighWaterMarkCounter* _memory_used_counter = nullptr;

    std::shared_ptr<ResourceContext> _resource_ctx;

    std::shared_ptr<QueryContext> _query_context;

    // Fragment and node id of the destination exchange node this receiver is used by.

    TUniqueId _fragment_instance_id;

    PlanNodeId _dest_node_id;

    // Row schema, copied from the caller of CreateRecvr().

    RowDescriptor _row_desc;

    // True if this reciver merges incoming rows from different senders. Per-sender

    // row batch queues are maintained in this case.

    bool _is_merging;

    bool _is_closed;

    std::unique_ptr<MemTracker> _mem_tracker;

    // Managed by object pool

    std::vector<SenderQueue*> _sender_queues;

    std::atomic<size_t> _sender_queue_mem_limit;

    std::unique_ptr<VSortedRunMerger> _merger;

    ObjectPool _sender_queue_pool;

    RuntimeProfile* _profile = nullptr;

    RuntimeProfile::Counter* _remote_bytes_received_counter = nullptr;

    RuntimeProfile::Counter* _local_bytes_received_counter = nullptr;

    RuntimeProfile::Counter* _deserialize_row_batch_timer = nullptr;

    RuntimeProfile::Counter* _first_batch_wait_total_timer = nullptr;

    RuntimeProfile::Counter* _buffer_full_total_timer = nullptr;

    RuntimeProfile::Counter* _data_arrival_timer = nullptr;

    RuntimeProfile::Counter* _decompress_timer = nullptr;

    RuntimeProfile::Counter* _decompress_bytes = nullptr;

    // Number of blocks received

    RuntimeProfile::Counter* _blocks_produced_counter = nullptr;

    RuntimeProfile::Counter* _max_wait_worker_time = nullptr;

    RuntimeProfile::Counter* _max_wait_to_process_time = nullptr;

    RuntimeProfile::Counter* _max_find_recvr_time = nullptr;

    std::vector<std::shared_ptr<Dependency>> _sender_to_local_channel_dependency;

};

class VDataStreamRecvr::SenderQueue {

public:

    SenderQueue(VDataStreamRecvr* parent_recvr, int num_senders,

                std::shared_ptr<Dependency> local_channel_dependency);

    ~SenderQueue();

    Status get_batch(Block* next_block, bool* eos);

    bool has_data_or_finished();

    Status add_block(std::unique_ptr<PBlock> pblock, int be_number, int64_t packet_seq,

                     ::google::protobuf::Closure** done, const int64_t wait_for_worker,

                     const uint64_t time_to_find_recvr);

    Status add_blocks(const PTransmitDataParams* request, ::google::protobuf::Closure** done,

                      const int64_t wait_for_worker, const uint64_t time_to_find_recvr);

    std::string debug_string();

    void add_block(Block* block, bool use_move);

    void decrement_senders(int sender_id);

    void cancel(Status cancel_status);

    void close();

    void set_dependency(std::shared_ptr<Dependency> dependency) {

        // Assign under _lock: set_source_ready() (reached via decrement_senders/cancel/close

        // on other threads) reads _source_dependency while holding _lock, so a lock-free

        // shared_ptr assignment here would race with that read.

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

        _source_dependency = dependency;

        // A queue created with zero senders (bucket-shuffle orphan instance, see

        // ExchangeLocalState::create_stream_recvr) never goes through decrement_senders,

        // so the usual reached-zero set_ready never fires — mark it ready at wiring time

        // or its task blocks forever on SHUFFLE_DATA_DEPENDENCY.

        if (_num_remaining_senders == 0) {

            set_source_ready(l);

        }

    }

protected:

    struct BlockItem;

    void add_blocks_memory_usage(int64_t size);

    void sub_blocks_memory_usage(int64_t size);

    bool exceeds_limit();

    friend class ExchangeLocalState;

    void set_source_ready(std::lock_guard<std::mutex>&);

    void run_block_queue_done_callbacks(std::list<BlockItem>& block_queue);

    // Not managed by this class

    VDataStreamRecvr* _recvr = nullptr;

    std::mutex _lock;

    bool _is_cancelled;

    Status _cancel_status;

    int _num_remaining_senders;

    std::unique_ptr<MemTracker> _queue_mem_tracker;

    // `BlockItem` is used in `_block_queue` to handle both local and remote exchange blocks.

    // For local exchange blocks, `BlockUPtr` is used directly without any modification.

    // For remote exchange blocks, the `pblock` is stored in `BlockItem`.

    // When `getBlock` is called, the `pblock` is deserialized into a usable block.

    struct BlockItem {

        Status get_block(BlockUPtr& block) {

            if (!_block) {

                DCHECK(_pblock);

                SCOPED_RAW_TIMER(&_deserialize_time);

                _block = Block::create_unique();

                RETURN_IF_ERROR_OR_CATCH_EXCEPTION(

                        _block->deserialize(*_pblock, &_decompress_bytes, &_decompress_time));

            }

            block.swap(_block);

            _block.reset();

            return Status::OK();

        }

        size_t block_byte_size() const { return _block_byte_size; }

        int64_t deserialize_time() const { return _deserialize_time; }

        int64_t decompress_time() const { return _decompress_time; }

        size_t decompress_bytes() const { return _decompress_bytes; }

        BlockItem();

        BlockItem(BlockUPtr&& block, size_t block_byte_size);

        BlockItem(std::unique_ptr<PBlock>&& pblock, size_t block_byte_size);

        BlockItem(BlockItem&&) noexcept;

        BlockItem& operator=(BlockItem&&) noexcept;

        ~BlockItem();

        void set_done(google::protobuf::Closure* done) {

            // The done callback is only set when the queue memory limit is exceeded.

            _done_cb = done;

            _wait_timer.start();

        }

        void call_done(VDataStreamRecvr* recvr) {

            if (_done_cb != nullptr) {

                _done_cb->Run();

                _done_cb = nullptr;

                _wait_timer.stop();

                int64_t elapse_time = _wait_timer.elapsed_time();

                if (recvr->_max_wait_to_process_time->value() < elapse_time) {

                    recvr->_max_wait_to_process_time->set(elapse_time);

                }

                recvr->_buffer_full_total_timer->update(elapse_time);

            }

        }

    private:

        BlockUPtr _block;

        std::unique_ptr<PBlock> _pblock;

        size_t _block_byte_size = 0;

        int64_t _deserialize_time = 0;

        int64_t _decompress_time = 0;

        size_t _decompress_bytes = 0;

        google::protobuf::Closure* _done_cb = nullptr;

        MonotonicStopWatch _wait_timer;

    };

    std::list<BlockItem> _block_queue;

    // sender_id

    std::unordered_set<int> _sender_eos_set;

    // be_number => packet_seq

    std::unordered_map<int, int64_t> _packet_seq_map;

    std::shared_ptr<Dependency> _source_dependency;

    std::shared_ptr<Dependency> _local_channel_dependency;

};

} // namespace doris