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

// bthread - An M:N threading library to make applications more concurrent.

// Date: Tue Jul 10 17:40:58 CST 2012

#ifndef BTHREAD_TASK_GROUP_H

#define BTHREAD_TASK_GROUP_H

#include "butil/time.h"                             // cpuwide_time_ns

#include "bthread/task_control.h"

#include "bthread/task_meta.h"                     // bthread_t, TaskMeta

#include "bthread/work_stealing_queue.h"           // WorkStealingQueue

#include "bthread/remote_task_queue.h"             // RemoteTaskQueue

#include "butil/resource_pool.h"                    // ResourceId

#include "bthread/parking_lot.h"

#include "bthread/prime_offset.h"

namespace bthread {

// For exiting a bthread.

class ExitException : public std::exception {

public:

    explicit ExitException(void* value) : _value(value) {}

    ~ExitException() throw() {}

    const char* what() const throw() override {

        return "ExitException";

    }

    void* value() const {

        return _value;

    }

private:

    void* _value;

};

// Refer to https://rigtorp.se/isatomic/, On the modern CPU microarchitectures

// (Skylake and Zen 2) AVX/AVX2 128b/256b aligned loads and stores are atomic

// even though Intel and AMD officially doesn’t guarantee this.

// On X86, SSE instructions can ensure atomic loads and stores.

// Starting from Armv8.4-A, neon can ensure atomic loads and stores.

// Otherwise, use mutex to guarantee atomicity.

class AtomicInteger128 {

public:

    struct BAIDU_CACHELINE_ALIGNMENT Value {

        int64_t v1;

        int64_t v2;

    };

    AtomicInteger128() = default;

    explicit AtomicInteger128(Value value) : _value(value) {}

    Value load() const;

    Value load_unsafe() const {

        return _value;

    }

    void store(Value value);

private:

    Value _value{};

    // Used to protect `_cpu_time_stat' when __x86_64__ and __ARM_NEON is not defined.

    FastPthreadMutex _mutex;

};

// Thread-local group of tasks.

// Notice that most methods involving context switching are static otherwise

// pointer `this' may change after wakeup. The **pg parameters in following

// function are updated before returning.

class TaskGroup {

public:

    // Create task `fn(arg)' with attributes `attr' in TaskGroup *pg and put

    // the identifier into `tid'. Switch to the new task and schedule old task

    // to run.

    // Return 0 on success, errno otherwise.

    static int start_foreground(TaskGroup** pg,

                                bthread_t* __restrict tid,

                                const bthread_attr_t* __restrict attr,

                                void * (*fn)(void*),

                                void* __restrict arg);

    // Create task `fn(arg)' with attributes `attr' in this TaskGroup, put the

    // identifier into `tid'. Schedule the new thread to run.

    //   Called from worker: start_background<false>

    //   Called from non-worker: start_background<true>

    // Return 0 on success, errno otherwise.

    template <bool REMOTE>

    int start_background(bthread_t* __restrict tid,

                         const bthread_attr_t* __restrict attr,

                         void * (*fn)(void*),

                         void* __restrict arg);

    // Suspend caller and run next bthread in TaskGroup *pg.

    static void sched(TaskGroup** pg);

    static void ending_sched(TaskGroup** pg);

    // Suspend caller and run bthread `next_tid' in TaskGroup *pg.

    // Purpose of this function is to avoid pushing `next_tid' to _rq and

    // then being popped by sched(pg), which is not necessary.

    static void sched_to(TaskGroup** pg, TaskMeta* next_meta, bool cur_ending);

    static void sched_to(TaskGroup** pg, bthread_t next_tid);

    static void exchange(TaskGroup** pg, TaskMeta* next_meta);

    // The callback will be run in the beginning of next-run bthread.

    // Can't be called by current bthread directly because it often needs

    // the target to be suspended already.

    typedef void (*RemainedFn)(void*);

    void set_remained(RemainedFn cb, void* arg) {

        _last_context_remained = cb;

        _last_context_remained_arg = arg;

    }

    // Suspend caller for at least |timeout_us| microseconds.

    // If |timeout_us| is 0, this function does nothing.

    // If |group| is NULL or current thread is non-bthread, call usleep(3)

    // instead. This function does not create thread-local TaskGroup.

    // Returns: 0 on success, -1 otherwise and errno is set.

    static int usleep(TaskGroup** pg, uint64_t timeout_us);

    // Suspend caller and run another bthread. When the caller will resume

    // is undefined.

    static void yield(TaskGroup** pg);

    // Suspend caller until bthread `tid' terminates.

    static int join(bthread_t tid, void** return_value);

    // Returns true iff the bthread `tid' still exists. Notice that it is

    // just the result at this very moment which may change soon.

    // Don't use this function unless you have to. Never write code like this:

    //    if (exists(tid)) {

    //        Wait for events of the thread.   // Racy, may block indefinitely.

    //    }

    static bool exists(bthread_t tid);

    // Put attribute associated with `tid' into `*attr'.

    // Returns 0 on success, -1 otherwise and errno is set.

    static int get_attr(bthread_t tid, bthread_attr_t* attr);

    // Get/set TaskMeta.stop of the tid.

    static void set_stopped(bthread_t tid);

    static bool is_stopped(bthread_t tid);

    // The bthread running run_main_task();

    bthread_t main_tid() const { return _main_tid; }

    TaskStatistics main_stat() const;

    // Routine of the main task which should be called from a dedicated pthread.

    void run_main_task();

    // current_task is a function in macOS 10.0+

#ifdef current_task

#undef current_task

#endif

    // Meta/Identifier of current task in this group.

    TaskMeta* current_task() const { return _cur_meta; }

    bthread_t current_tid() const { return _cur_meta->tid; }

    // Uptime of current task in nanoseconds.

    int64_t current_uptime_ns() const

    { return butil::cpuwide_time_ns() - _cur_meta->cpuwide_start_ns; }

    // True iff current task is the one running run_main_task()

    bool is_current_main_task() const { return current_tid() == _main_tid; }

    // True iff current task is in pthread-mode.

    bool is_current_pthread_task() const

    { return _cur_meta->stack == _main_stack; }

    // Active time in nanoseconds spent by this TaskGroup.

    int64_t cumulated_cputime_ns() const;

    // Push a bthread into the runqueue

    void ready_to_run(TaskMeta* meta, bool nosignal = false);

    // Flush tasks pushed to rq but signalled.

    void flush_nosignal_tasks();

    // Push a bthread into the runqueue from another non-worker thread.

    void ready_to_run_remote(TaskMeta* meta, bool nosignal = false);

    void flush_nosignal_tasks_remote_locked(butil::Mutex& locked_mutex);

    void flush_nosignal_tasks_remote();

    // Automatically decide the caller is remote or local, and call

    // the corresponding function.

    void ready_to_run_general(TaskMeta* meta, bool nosignal = false);

    void flush_nosignal_tasks_general();

    // The TaskControl that this TaskGroup belongs to.

    TaskControl* control() const { return _control; }

    // Call this instead of delete.

    void destroy_self();

    // Wake up blocking ops in the thread.

    // Returns 0 on success, errno otherwise.

    static int interrupt(bthread_t tid, TaskControl* c, bthread_tag_t tag);

    // Get the meta associate with the task.

    static TaskMeta* address_meta(bthread_t tid);

    // Push a task into _rq, if _rq is full, retry after some time. This

    // process make go on indefinitely.

    void push_rq(bthread_t tid);

    // Returns size of local run queue.

    size_t rq_size() const {

        return _rq.volatile_size();

    }

    bthread_tag_t tag() const { return _tag; }

    pthread_t tid() const { return _tid; }

    int64_t current_task_cpu_clock_ns() {

        if (_last_cpu_clock_ns == 0) {

            return 0;

        }

        int64_t total_ns = _cur_meta->stat.cpu_usage_ns;

        total_ns += butil::cputhread_time_ns() - _last_cpu_clock_ns;

        return total_ns;

    }

private:

friend class TaskControl;

    // Last scheduling time, task type and cumulated CPU time.

    class CPUTimeStat {

        static constexpr int64_t LAST_SCHEDULING_TIME_MASK = 0x7FFFFFFFFFFFFFFFLL;

        static constexpr int64_t TASK_TYPE_MASK = 0x8000000000000000LL;

    public:

        CPUTimeStat() : _last_run_ns_and_type(0), _cumulated_cputime_ns(0) {}

        CPUTimeStat(AtomicInteger128::Value value)

            : _last_run_ns_and_type(value.v1), _cumulated_cputime_ns(value.v2) {}

        // Convert to AtomicInteger128::Value for atomic operations.

        explicit operator AtomicInteger128::Value() const {

            return {_last_run_ns_and_type, _cumulated_cputime_ns};

        }

        void set_last_run_ns(int64_t last_run_ns, bool main_task) {

            _last_run_ns_and_type = (last_run_ns & LAST_SCHEDULING_TIME_MASK) |

                                    (static_cast<int64_t>(main_task) << 63);

        }

        int64_t last_run_ns() const {

            return _last_run_ns_and_type & LAST_SCHEDULING_TIME_MASK;

        }

        int64_t last_run_ns_and_type() const {

            return _last_run_ns_and_type;

        }

        bool is_main_task() const {

            return _last_run_ns_and_type & TASK_TYPE_MASK;

        }

        void add_cumulated_cputime_ns(int64_t cputime_ns, bool main_task) {

            if (main_task) {

                return;

            }

            _cumulated_cputime_ns += cputime_ns;

        }

        int64_t cumulated_cputime_ns() const {

            return _cumulated_cputime_ns;

        }

    private:

        // The higher bit for task type, main task is 1, otherwise 0.

        // Lowest 63 bits for last scheduling time.

        int64_t _last_run_ns_and_type;

        // Cumulated CPU time in nanoseconds.

        int64_t _cumulated_cputime_ns;

    };

    class AtomicCPUTimeStat {

    public:

        CPUTimeStat load() const {

            return  _cpu_time_stat.load();

        }

        CPUTimeStat load_unsafe() const {

            return _cpu_time_stat.load_unsafe();

        }

        void store(CPUTimeStat cpu_time_stat) {

            _cpu_time_stat.store(AtomicInteger128::Value(cpu_time_stat));

        }

    private:

        AtomicInteger128 _cpu_time_stat;

    };

    // You shall use TaskControl::create_group to create new instance.

    explicit TaskGroup(TaskControl* c);

    int init(size_t runqueue_capacity);

    // You shall call destroy_selfm() instead of destructor because deletion

    // of groups are postponed to avoid race.

    ~TaskGroup();

#ifdef BUTIL_USE_ASAN

    static void asan_task_runner(intptr_t);

#endif // BUTIL_USE_ASAN

    static void task_runner(intptr_t skip_remained);

    // Callbacks for set_remained()

    static void _release_last_context(void*);

    static void _add_sleep_event(void*);

    struct ReadyToRunArgs {

        bthread_tag_t tag;

        TaskMeta* meta;

        bool nosignal;

    };

    static void ready_to_run_in_worker(void*);

    static void ready_to_run_in_worker_ignoresignal(void*);

    static void priority_to_run(void*);

    // Wait for a task to run.

    // Returns true on success, false is treated as permanent error and the

    // loop calling this function should end.

    bool wait_task(bthread_t* tid);

    bool steal_task(bthread_t* tid) {

        if (_remote_rq.pop(tid)) {

            return true;

        }

#ifndef BTHREAD_DONT_SAVE_PARKING_STATE

        _last_pl_state = _pl->get_state();

#endif

        return _control->steal_task(tid, &_steal_seed, _steal_offset);

    }

    void set_tag(bthread_tag_t tag) { _tag = tag; }

    void set_pl(ParkingLot* pl) { _pl = pl; }

    static bool is_main_task(TaskGroup* g, bthread_t tid) {

        return g->_main_tid == tid;

    }

    TaskMeta* _cur_meta{NULL};

    // the control that this group belongs to

    TaskControl* _control{NULL};

    int _num_nosignal{0};

    int _nsignaled{0};

    AtomicCPUTimeStat _cpu_time_stat;

    // last thread cpu clock

    int64_t _last_cpu_clock_ns{0};

    size_t _nswitch{0};

    RemainedFn _last_context_remained{NULL};

    void* _last_context_remained_arg{NULL};

    ParkingLot* _pl{NULL};

#ifndef BTHREAD_DONT_SAVE_PARKING_STATE

    ParkingLot::State _last_pl_state;

#endif

    size_t _steal_seed{butil::fast_rand()};

    size_t _steal_offset{prime_offset(_steal_seed)};

    ContextualStack* _main_stack{NULL};

    bthread_t _main_tid{INVALID_BTHREAD};

    WorkStealingQueue<bthread_t> _rq;

    RemoteTaskQueue _remote_rq;

    int _remote_num_nosignal{0};

    int _remote_nsignaled{0};

    int _sched_recursive_guard{0};

    // tag of this taskgroup

    bthread_tag_t _tag{BTHREAD_TAG_DEFAULT};

    // Worker thread id.

    pthread_t _tid{};

};

}  // namespace bthread

#include "task_group_inl.h"

#endif  // BTHREAD_TASK_GROUP_H