// 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/threadpool.cc
// and modified by Doris

#include "util/threadpool.h"

#include <algorithm>
#include <cstdint>
#include <limits>
#include <ostream>
#include <thread>
#include <utility>

#include "common/logging.h"
#include "gutil/map-util.h"
#include "gutil/port.h"
#include "gutil/strings/substitute.h"
#include "util/debug/sanitizer_scopes.h"
#include "util/scoped_cleanup.h"
#include "util/thread.h"

namespace doris {
using namespace ErrorCode;

using std::string;
using strings::Substitute;

class FunctionRunnable : public Runnable {
public:
    explicit FunctionRunnable(std::function<void()> func) : _func(std::move(func)) {}

    void run() override { _func(); }

private:
    std::function<void()> _func;
};

ThreadPoolBuilder::ThreadPoolBuilder(string name)
        : _name(std::move(name)),
          _min_threads(0),
          _max_threads(std::thread::hardware_concurrency()),
          _max_queue_size(std::numeric_limits<int>::max()),
          _idle_timeout(std::chrono::milliseconds(500)) {}

ThreadPoolBuilder& ThreadPoolBuilder::set_min_threads(int min_threads) {
    CHECK_GE(min_threads, 0);
    _min_threads = min_threads;
    return *this;
}

ThreadPoolBuilder& ThreadPoolBuilder::set_max_threads(int max_threads) {
    CHECK_GT(max_threads, 0);
    _max_threads = max_threads;
    return *this;
}

ThreadPoolBuilder& ThreadPoolBuilder::set_max_queue_size(int max_queue_size) {
    _max_queue_size = max_queue_size;
    return *this;
}

ThreadPoolBuilder& ThreadPoolBuilder::set_cgroup_cpu_ctl(
        std::weak_ptr<CgroupCpuCtl> cgroup_cpu_ctl) {
    _cgroup_cpu_ctl = cgroup_cpu_ctl;
    return *this;
}

ThreadPoolToken::ThreadPoolToken(ThreadPool* pool, ThreadPool::ExecutionMode mode,
                                 int max_concurrency)
        : _mode(mode),
          _pool(pool),
          _state(State::IDLE),
          _active_threads(0),
          _max_concurrency(max_concurrency),
          _num_submitted_tasks(0),
          _num_unsubmitted_tasks(0) {
    if (max_concurrency == 1 && mode != ThreadPool::ExecutionMode::SERIAL) {
        _mode = ThreadPool::ExecutionMode::SERIAL;
    }
}

ThreadPoolToken::~ThreadPoolToken() {
    shutdown();
    _pool->release_token(this);
}

Status ThreadPoolToken::submit(std::shared_ptr<Runnable> r) {
    return _pool->do_submit(std::move(r), this);
}

Status ThreadPoolToken::submit_func(std::function<void()> f) {
    return submit(std::make_shared<FunctionRunnable>(std::move(f)));
}

void ThreadPoolToken::shutdown() {
    std::unique_lock<std::mutex> l(_pool->_lock);
    _pool->check_not_pool_thread_unlocked();

    // Clear the queue under the lock, but defer the releasing of the tasks
    // outside the lock, in case there are concurrent threads wanting to access
    // the ThreadPool. The task's destructors may acquire locks, etc, so this
    // also prevents lock inversions.
    std::deque<ThreadPool::Task> to_release = std::move(_entries);
    _pool->_total_queued_tasks -= to_release.size();

    switch (state()) {
    case State::IDLE:
        // There were no tasks outstanding; we can quiesce the token immediately.
        transition(State::QUIESCED);
        break;
    case State::RUNNING:
        // There were outstanding tasks. If any are still running, switch to
        // QUIESCING and wait for them to finish (the worker thread executing
        // the token's last task will switch the token to QUIESCED). Otherwise,
        // we can quiesce the token immediately.

        // Note: this is an O(n) operation, but it's expected to be infrequent.
        // Plus doing it this way (rather than switching to QUIESCING and waiting
        // for a worker thread to process the queue entry) helps retain state
        // transition symmetry with ThreadPool::shutdown.
        for (auto it = _pool->_queue.begin(); it != _pool->_queue.end();) {
            if (*it == this) {
                it = _pool->_queue.erase(it);
            } else {
                it++;
            }
        }

        if (_active_threads == 0) {
            transition(State::QUIESCED);
            break;
        }
        transition(State::QUIESCING);
        [[fallthrough]];
    case State::QUIESCING:
        // The token is already quiescing. Just wait for a worker thread to
        // switch it to QUIESCED.
        _not_running_cond.wait(l, [this]() { return state() == State::QUIESCED; });
        break;
    default:
        break;
    }
}

void ThreadPoolToken::wait() {
    std::unique_lock<std::mutex> l(_pool->_lock);
    _pool->check_not_pool_thread_unlocked();
    _not_running_cond.wait(l, [this]() { return !is_active(); });
}

void ThreadPoolToken::transition(State new_state) {
#ifndef NDEBUG
    CHECK_NE(_state, new_state);

    switch (_state) {
    case State::IDLE:
        CHECK(new_state == State::RUNNING || new_state == State::QUIESCED);
        if (new_state == State::RUNNING) {
            CHECK(!_entries.empty());
        } else {
            CHECK(_entries.empty());
            CHECK_EQ(_active_threads, 0);
        }
        break;
    case State::RUNNING:
        CHECK(new_state == State::IDLE || new_state == State::QUIESCING ||
              new_state == State::QUIESCED);
        CHECK(_entries.empty());
        if (new_state == State::QUIESCING) {
            CHECK_GT(_active_threads, 0);
        }
        break;
    case State::QUIESCING:
        CHECK(new_state == State::QUIESCED);
        CHECK_EQ(_active_threads, 0);
        break;
    case State::QUIESCED:
        CHECK(false); // QUIESCED is a terminal state
        break;
    default:
        LOG(FATAL) << "Unknown token state: " << _state;
    }
#endif

    // Take actions based on the state we're entering.
    switch (new_state) {
    case State::IDLE:
    case State::QUIESCED:
        _not_running_cond.notify_all();
        break;
    default:
        break;
    }

    _state = new_state;
}

const char* ThreadPoolToken::state_to_string(State s) {
    switch (s) {
    case State::IDLE:
        return "IDLE";
        break;
    case State::RUNNING:
        return "RUNNING";
        break;
    case State::QUIESCING:
        return "QUIESCING";
        break;
    case State::QUIESCED:
        return "QUIESCED";
        break;
    }
    return "<cannot reach here>";
}

bool ThreadPoolToken::need_dispatch() {
    return _state == ThreadPoolToken::State::IDLE ||
           (_mode == ThreadPool::ExecutionMode::CONCURRENT &&
            _num_submitted_tasks < _max_concurrency);
}

ThreadPool::ThreadPool(const ThreadPoolBuilder& builder)
        : _name(builder._name),
          _min_threads(builder._min_threads),
          _max_threads(builder._max_threads),
          _max_queue_size(builder._max_queue_size),
          _idle_timeout(builder._idle_timeout),
          _pool_status(Status::Uninitialized("The pool was not initialized.")),
          _num_threads(0),
          _num_threads_pending_start(0),
          _active_threads(0),
          _total_queued_tasks(0),
          _cgroup_cpu_ctl(builder._cgroup_cpu_ctl),
          _tokenless(new_token(ExecutionMode::CONCURRENT)) {}

ThreadPool::~ThreadPool() {
    // There should only be one live token: the one used in tokenless submission.
    CHECK_EQ(1, _tokens.size()) << strings::Substitute(
            "Threadpool $0 destroyed with $1 allocated tokens", _name, _tokens.size());
    shutdown();
}

Status ThreadPool::init() {
    if (!_pool_status.is<UNINITIALIZED>()) {
        return Status::NotSupported("The thread pool {} is already initialized", _name);
    }
    _pool_status = Status::OK();
    _num_threads_pending_start = _min_threads;
    for (int i = 0; i < _min_threads; i++) {
        Status status = create_thread();
        if (!status.ok()) {
            shutdown();
            return status;
        }
    }
    return Status::OK();
}

void ThreadPool::shutdown() {
    debug::ScopedTSANIgnoreReadsAndWrites ignore_tsan;
    std::unique_lock<std::mutex> l(_lock);
    check_not_pool_thread_unlocked();

    // Note: this is the same error seen at submission if the pool is at
    // capacity, so clients can't tell them apart. This isn't really a practical
    // concern though because shutting down a pool typically requires clients to
    // be quiesced first, so there's no danger of a client getting confused.
    // Not print stack trace here
    _pool_status = Status::Error<SERVICE_UNAVAILABLE, false>(
            "The thread pool {} has been shut down.", _name);

    // Clear the various queues under the lock, but defer the releasing
    // of the tasks outside the lock, in case there are concurrent threads
    // wanting to access the ThreadPool. The task's destructors may acquire
    // locks, etc, so this also prevents lock inversions.
    _queue.clear();

    std::deque<std::deque<Task>> to_release;
    for (auto* t : _tokens) {
        if (!t->_entries.empty()) {
            to_release.emplace_back(std::move(t->_entries));
        }
        switch (t->state()) {
        case ThreadPoolToken::State::IDLE:
            // The token is idle; we can quiesce it immediately.
            t->transition(ThreadPoolToken::State::QUIESCED);
            break;
        case ThreadPoolToken::State::RUNNING:
            // The token has tasks associated with it. If they're merely queued
            // (i.e. there are no active threads), the tasks will have been removed
            // above and we can quiesce immediately. Otherwise, we need to wait for
            // the threads to finish.
            t->transition(t->_active_threads > 0 ? ThreadPoolToken::State::QUIESCING
                                                 : ThreadPoolToken::State::QUIESCED);
            break;
        default:
            break;
        }
    }

    // The queues are empty. Wake any sleeping worker threads and wait for all
    // of them to exit. Some worker threads will exit immediately upon waking,
    // while others will exit after they finish executing an outstanding task.
    _total_queued_tasks = 0;
    while (!_idle_threads.empty()) {
        _idle_threads.front().not_empty.notify_one();
        _idle_threads.pop_front();
    }

    _no_threads_cond.wait(l, [this]() { return _num_threads + _num_threads_pending_start == 0; });

    // All the threads have exited. Check the state of each token.
    for (auto* t : _tokens) {
        DCHECK(t->state() == ThreadPoolToken::State::IDLE ||
               t->state() == ThreadPoolToken::State::QUIESCED);
    }
}

std::unique_ptr<ThreadPoolToken> ThreadPool::new_token(ExecutionMode mode, int max_concurrency) {
    std::lock_guard<std::mutex> l(_lock);
    std::unique_ptr<ThreadPoolToken> t(new ThreadPoolToken(this, mode, max_concurrency));
    InsertOrDie(&_tokens, t.get());
    return t;
}

void ThreadPool::release_token(ThreadPoolToken* t) {
    std::lock_guard<std::mutex> l(_lock);
    CHECK(!t->is_active()) << strings::Substitute("Token with state $0 may not be released",
                                                  ThreadPoolToken::state_to_string(t->state()));
    CHECK_EQ(1, _tokens.erase(t));
}

Status ThreadPool::submit(std::shared_ptr<Runnable> r) {
    return do_submit(std::move(r), _tokenless.get());
}

Status ThreadPool::submit_func(std::function<void()> f) {
    return submit(std::make_shared<FunctionRunnable>(std::move(f)));
}

Status ThreadPool::do_submit(std::shared_ptr<Runnable> r, ThreadPoolToken* token) {
    DCHECK(token);
    std::chrono::time_point<std::chrono::system_clock> submit_time =
            std::chrono::system_clock::now();

    std::unique_lock<std::mutex> l(_lock);
    if (PREDICT_FALSE(!_pool_status.ok())) {
        return _pool_status;
    }

    if (PREDICT_FALSE(!token->may_submit_new_tasks())) {
        return Status::Error<SERVICE_UNAVAILABLE>("Thread pool({}) token was shut down", _name);
    }

    // Size limit check.
    int64_t capacity_remaining = static_cast<int64_t>(_max_threads) - _active_threads +
                                 static_cast<int64_t>(_max_queue_size) - _total_queued_tasks;
    if (capacity_remaining < 1) {
        return Status::Error<SERVICE_UNAVAILABLE>(
                "Thread pool {} is at capacity ({}/{} tasks running, {}/{} tasks queued)", _name,
                _num_threads + _num_threads_pending_start, _max_threads, _total_queued_tasks,
                _max_queue_size);
    }

    // Should we create another thread?

    // We assume that each current inactive thread will grab one item from the
    // queue.  If it seems like we'll need another thread, we create one.
    //
    // Rather than creating the thread here, while holding the lock, we defer
    // it to down below. This is because thread creation can be rather slow
    // (hundreds of milliseconds in some cases) and we'd like to allow the
    // existing threads to continue to process tasks while we do so.
    //
    // In theory, a currently active thread could finish immediately after this
    // calculation but before our new worker starts running. This would mean we
    // created a thread we didn't really need. However, this race is unavoidable
    // and harmless.
    //
    // Of course, we never create more than _max_threads threads no matter what.
    int threads_from_this_submit =
            token->is_active() && token->mode() == ExecutionMode::SERIAL ? 0 : 1;
    int inactive_threads = _num_threads + _num_threads_pending_start - _active_threads;
    int additional_threads =
            static_cast<int>(_queue.size()) + threads_from_this_submit - inactive_threads;
    bool need_a_thread = false;
    if (additional_threads > 0 && _num_threads + _num_threads_pending_start < _max_threads) {
        need_a_thread = true;
        _num_threads_pending_start++;
    }

    Task task;
    task.runnable = std::move(r);
    task.submit_time = submit_time;

    // Add the task to the token's queue.
    ThreadPoolToken::State state = token->state();
    DCHECK(state == ThreadPoolToken::State::IDLE || state == ThreadPoolToken::State::RUNNING);
    token->_entries.emplace_back(std::move(task));
    // When we need to execute the task in the token, we submit the token object to the queue.
    // There are currently two places where tokens will be submitted to the queue:
    // 1. When submitting a new task, if the token is still in the IDLE state,
    //    or the concurrency of the token has not reached the online level, it will be added to the queue.
    // 2. When the dispatch thread finishes executing a task:
    //    1. If it is a SERIAL token, and there are unsubmitted tasks, submit them to the queue.
    //    2. If it is a CONCURRENT token, and there are still unsubmitted tasks, and the upper limit of concurrency is not reached,
    //       then submitted to the queue.
    if (token->need_dispatch()) {
        _queue.emplace_back(token);
        ++token->_num_submitted_tasks;
        if (state == ThreadPoolToken::State::IDLE) {
            token->transition(ThreadPoolToken::State::RUNNING);
        }
    } else {
        ++token->_num_unsubmitted_tasks;
    }
    _total_queued_tasks++;

    // Wake up an idle thread for this task. Choosing the thread at the front of
    // the list ensures LIFO semantics as idling threads are also added to the front.
    //
    // If there are no idle threads, the new task remains on the queue and is
    // processed by an active thread (or a thread we're about to create) at some
    // point in the future.
    if (!_idle_threads.empty()) {
        _idle_threads.front().not_empty.notify_one();
        _idle_threads.pop_front();
    }
    l.unlock();

    if (need_a_thread) {
        Status status = create_thread();
        if (!status.ok()) {
            l.lock();
            _num_threads_pending_start--;
            if (_num_threads + _num_threads_pending_start == 0) {
                // If we have no threads, we can't do any work.
                return status;
            }
            // If we failed to create a thread, but there are still some other
            // worker threads, log a warning message and continue.
            LOG(WARNING) << "Thread pool " << _name
                         << " failed to create thread: " << status.to_string();
        }
    }

    return Status::OK();
}

void ThreadPool::wait() {
    std::unique_lock<std::mutex> l(_lock);
    check_not_pool_thread_unlocked();
    _idle_cond.wait(l, [this]() { return _total_queued_tasks == 0 && _active_threads == 0; });
}

void ThreadPool::dispatch_thread() {
    std::unique_lock<std::mutex> l(_lock);
    debug::ScopedTSANIgnoreReadsAndWrites ignore_tsan;
    InsertOrDie(&_threads, Thread::current_thread());
    DCHECK_GT(_num_threads_pending_start, 0);
    _num_threads++;
    _num_threads_pending_start--;

    if (std::shared_ptr<CgroupCpuCtl> cg_cpu_ctl_sptr = _cgroup_cpu_ctl.lock()) {
        static_cast<void>(cg_cpu_ctl_sptr->add_thread_to_cgroup());
    }

    // Owned by this worker thread and added/removed from _idle_threads as needed.
    IdleThread me;

    while (true) {
        // Note: Status::Aborted() is used to indicate normal shutdown.
        if (!_pool_status.ok()) {
            VLOG_CRITICAL << "DispatchThread exiting: " << _pool_status.to_string();
            break;
        }

        if (_num_threads + _num_threads_pending_start > _max_threads) {
            break;
        }

        if (_queue.empty()) {
            // There's no work to do, let's go idle.
            //
            // Note: if FIFO behavior is desired, it's as simple as changing this to push_back().
            _idle_threads.push_front(me);
            SCOPED_CLEANUP({
                // For some wake ups (i.e. shutdown or do_submit) this thread is
                // guaranteed to be unlinked after being awakened. In others (i.e.
                // spurious wake-up or Wait timeout), it'll still be linked.
                if (me.is_linked()) {
                    _idle_threads.erase(_idle_threads.iterator_to(me));
                }
            });
            if (me.not_empty.wait_for(l, _idle_timeout) == std::cv_status::timeout) {
                // After much investigation, it appears that pthread condition variables have
                // a weird behavior in which they can return ETIMEDOUT from timed_wait even if
                // another thread did in fact signal. Apparently after a timeout there is some
                // brief period during which another thread may actually grab the internal mutex
                // protecting the state, signal, and release again before we get the mutex. So,
                // we'll recheck the empty queue case regardless.
                if (_queue.empty() && _num_threads + _num_threads_pending_start > _min_threads) {
                    VLOG_NOTICE << "Releasing worker thread from pool " << _name << " after "
                                << std::chrono::duration_cast<std::chrono::milliseconds>(
                                           _idle_timeout)
                                           .count()
                                << "ms of idle time.";
                    break;
                }
            }
            continue;
        }

        // Get the next token and task to execute.
        ThreadPoolToken* token = _queue.front();
        _queue.pop_front();
        DCHECK_EQ(ThreadPoolToken::State::RUNNING, token->state());
        DCHECK(!token->_entries.empty());
        Task task = std::move(token->_entries.front());
        token->_entries.pop_front();
        token->_active_threads++;
        --_total_queued_tasks;
        ++_active_threads;

        l.unlock();

        // Execute the task
        task.runnable->run();

        // Destruct the task while we do not hold the lock.
        //
        // The task's destructor may be expensive if it has a lot of bound
        // objects, and we don't want to block submission of the threadpool.
        // In the worst case, the destructor might even try to do something
        // with this threadpool, and produce a deadlock.
        task.runnable.reset();
        l.lock();

        // Possible states:
        // 1. The token was shut down while we ran its task. Transition to QUIESCED.
        // 2. The token has no more queued tasks. Transition back to IDLE.
        // 3. The token has more tasks. Requeue it and transition back to RUNNABLE.
        ThreadPoolToken::State state = token->state();
        DCHECK(state == ThreadPoolToken::State::RUNNING ||
               state == ThreadPoolToken::State::QUIESCING);
        --token->_active_threads;
        --token->_num_submitted_tasks;

        // handle shutdown && idle
        if (token->_active_threads == 0) {
            if (state == ThreadPoolToken::State::QUIESCING) {
                DCHECK(token->_entries.empty());
                token->transition(ThreadPoolToken::State::QUIESCED);
            } else if (token->_entries.empty()) {
                token->transition(ThreadPoolToken::State::IDLE);
            }
        }

        // We decrease _num_submitted_tasks holding lock, so the following DCHECK works.
        DCHECK(token->_num_submitted_tasks < token->_max_concurrency);

        // If token->state is running and there are unsubmitted tasks in the token, we put
        // the token back.
        if (token->_num_unsubmitted_tasks > 0 && state == ThreadPoolToken::State::RUNNING) {
            // SERIAL: if _entries is not empty, then num_unsubmitted_tasks must be greater than 0.
            // CONCURRENT: we have to check _num_unsubmitted_tasks because there may be at least 2
            // threads are running for the token.
            _queue.emplace_back(token);
            ++token->_num_submitted_tasks;
            --token->_num_unsubmitted_tasks;
        }

        if (--_active_threads == 0) {
            _idle_cond.notify_all();
        }
    }

    // It's important that we hold the lock between exiting the loop and dropping
    // _num_threads. Otherwise it's possible someone else could come along here
    // and add a new task just as the last running thread is about to exit.
    CHECK(l.owns_lock());

    CHECK_EQ(_threads.erase(Thread::current_thread()), 1);
    _num_threads--;
    if (_num_threads + _num_threads_pending_start == 0) {
        _no_threads_cond.notify_all();

        // Sanity check: if we're the last thread exiting, the queue ought to be
        // empty. Otherwise it will never get processed.
        CHECK(_queue.empty());
        DCHECK_EQ(0, _total_queued_tasks);
    }
}

Status ThreadPool::create_thread() {
    return Thread::create("thread pool", strings::Substitute("$0 [worker]", _name),
                          &ThreadPool::dispatch_thread, this, nullptr);
}

void ThreadPool::check_not_pool_thread_unlocked() {
    Thread* current = Thread::current_thread();
    if (ContainsKey(_threads, current)) {
        LOG(FATAL) << strings::Substitute(
                "Thread belonging to thread pool '$0' with "
                "name '$1' called pool function that would result in deadlock",
                _name, current->name());
    }
}

Status ThreadPool::set_min_threads(int min_threads) {
    std::lock_guard<std::mutex> l(_lock);
    if (min_threads > _max_threads) {
        // min threads can not be set greater than max threads
        return Status::InternalError("set thread pool {} min_threads failed", _name);
    }
    _min_threads = min_threads;
    if (min_threads > _num_threads + _num_threads_pending_start) {
        int addition_threads = min_threads - _num_threads - _num_threads_pending_start;
        _num_threads_pending_start += addition_threads;
        for (int i = 0; i < addition_threads; i++) {
            Status status = create_thread();
            if (!status.ok()) {
                _num_threads_pending_start--;
                LOG(WARNING) << "Thread pool " << _name
                             << " failed to create thread: " << status.to_string();
                return status;
            }
        }
    }
    return Status::OK();
}

Status ThreadPool::set_max_threads(int max_threads) {
    std::lock_guard<std::mutex> l(_lock);
    if (_min_threads > max_threads) {
        // max threads can not be set less than min threads
        return Status::InternalError("set thread pool {} max_threads failed", _name);
    }

    _max_threads = max_threads;
    if (_max_threads > _num_threads + _num_threads_pending_start) {
        int addition_threads = _max_threads - _num_threads - _num_threads_pending_start;
        addition_threads = std::min(addition_threads, _total_queued_tasks);
        _num_threads_pending_start += addition_threads;
        for (int i = 0; i < addition_threads; i++) {
            Status status = create_thread();
            if (!status.ok()) {
                _num_threads_pending_start--;
                LOG(WARNING) << "Thread pool " << _name
                             << " failed to create thread: " << status.to_string();
                return status;
            }
        }
    }
    return Status::OK();
}

std::ostream& operator<<(std::ostream& o, ThreadPoolToken::State s) {
    return o << ThreadPoolToken::state_to_string(s);
}

} // namespace doris