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

// This file is copied from

// https://github.com/apache/impala/blob/branch-2.9.0/be/src/util/blocking-priority-queue.hpp

// and modified by Doris

#pragma once

#include <unistd.h>

#include <cassert>

#include <condition_variable>

#include <cstdint>

#include <mutex>

#include <queue>

#include "common/config.h"

#include "util/stopwatch.hpp"

namespace doris {

// Fixed capacity FIFO queue, where both blocking_get and blocking_put operations block

// if the queue is empty or full, respectively.

template <typename T>

class BlockingPriorityQueue {

public:

BlockingPriorityQueue(uint32_t max_elements)

: _shutdown(false),

_max_element(max_elements),

_upgrade_counter(0),

_total_get_wait_time(0),

_total_put_wait_time(0),

_get_waiting(0),

_put_waiting(0) {}

// Get an element from the queue, waiting indefinitely (or until timeout) for one to become available.

// Returns false if we were shut down prior to getting the element, and there

// are no more elements available.

// -- timeout_ms: 0 means wait indefinitely

bool blocking_get(T* out, uint32_t timeout_ms = 0) {

MonotonicStopWatch timer;

timer.start();

std::unique_lock unique_lock(_lock);

bool wait_successful = false;

if (timeout_ms > 0) {

while (!(_shutdown || !_queue.empty())) {

++_get_waiting;

if (_get_cv.wait_for(unique_lock, std::chrono::milliseconds(timeout_ms)) ==

std::cv_status::timeout) {

// timeout

wait_successful = _shutdown || !_queue.empty();

break;

}

}

} else {

while (!(_shutdown || !_queue.empty())) {

++_get_waiting;

_get_cv.wait(unique_lock);

}

wait_successful = true;

}

_total_get_wait_time += timer.elapsed_time();

if (wait_successful) {

if (_upgrade_counter > config::priority_queue_remaining_tasks_increased_frequency) {

std::priority_queue<T> tmp_queue;

while (!_queue.empty()) {

T v = _queue.top();

_queue.pop();

++v;

tmp_queue.push(v);

}

swap(_queue, tmp_queue);

_upgrade_counter = 0;

}

if (!_queue.empty()) {

*out = _queue.top();

_queue.pop();

++_upgrade_counter;

if (_put_waiting > 0) {

--_put_waiting;

unique_lock.unlock();

_put_cv.notify_one();

}

return true;

} else {

assert(_shutdown);

return false;

}

} else {

//time out

assert(!_shutdown);

return false;

}

}

bool non_blocking_get(T* out) {

MonotonicStopWatch timer;

timer.start();

std::unique_lock unique_lock(_lock);

if (!_queue.empty()) {

// 定期提高队列中残留的任务优先级

// 保证优先级较低的大查询不至于完全饿死

if (_upgrade_counter > config::priority_queue_remaining_tasks_increased_frequency) {

std::priority_queue<T> tmp_queue;

while (!_queue.empty()) {

T v = _queue.top();

_queue.pop();

++v;

tmp_queue.push(v);

}

swap(_queue, tmp_queue);

_upgrade_counter = 0;

}

*out = _queue.top();

_queue.pop();

++_upgrade_counter;

_total_get_wait_time += timer.elapsed_time();

if (_put_waiting > 0) {

--_put_waiting;

unique_lock.unlock();

_put_cv.notify_one();

}

return true;

}

return false;

}

// Puts an element into the queue, waiting indefinitely until there is space.

// If the queue is shut down, returns false.

bool blocking_put(const T& val) {

MonotonicStopWatch timer;

timer.start();

std::unique_lock unique_lock(_lock);

while (!(_shutdown || _queue.size() < _max_element)) {

++_put_waiting;

_put_cv.wait(unique_lock);

}

_total_put_wait_time += timer.elapsed_time();

if (_shutdown) {

return false;

}

_queue.push(val);

if (_get_waiting > 0) {

--_get_waiting;

unique_lock.unlock();

_get_cv.notify_one();

}

return true;

}

// Return false if queue full or has been shutdown.

bool try_put(const T& val) {

std::unique_lock unique_lock(_lock);

if (_queue.size() < _max_element && !_shutdown) {

_queue.push(val);

if (_get_waiting > 0) {

--_get_waiting;

unique_lock.unlock();

_get_cv.notify_one();

}

return true;

}

return false;

}

// Shut down the queue. Wakes up all threads waiting on blocking_get or blocking_put.

void shutdown() {

{

std::lock_guard l(_lock);

_shutdown = true;

}

_get_cv.notify_all();

_put_cv.notify_all();

}

uint32_t get_size() const {

std::lock_guard l(_lock);

return static_cast<uint32_t>(_queue.size());

}

uint32_t get_capacity() const { return _max_element; }

// Returns the total amount of time threads have blocked in blocking_get.

uint64_t total_get_wait_time() const { return _total_get_wait_time; }

// Returns the total amount of time threads have blocked in blocking_put.

uint64_t total_put_wait_time() const { return _total_put_wait_time; }

private:

bool _shutdown;

const int _max_element;

std::condition_variable _get_cv; // 'get' callers wait on this

std::condition_variable _put_cv; // 'put' callers wait on this

// _lock guards access to _queue, total_get_wait_time, and total_put_wait_time

mutable std::mutex _lock;

std::priority_queue<T> _queue;

int _upgrade_counter;

std::atomic<uint64_t> _total_get_wait_time;

std::atomic<uint64_t> _total_put_wait_time;

size_t _get_waiting;

size_t _put_waiting;

};

} // namespace doris