// Agora Media SDK

//

// Created by Yaqi Li in 2021-03.

// Copyright (c) 2021 Agora IO. All rights reserved.

//

#pragma once

#include <thread>

#include <list>

#include <string>

#include <functional>

#include <mutex>

#include <condition_variable>

#include <unordered_map>

#include <vector>

#include <future>

namespace agora {

namespace extension {

class ThreadPool {

public:

using FuncType = std::function<void()>;

explicit ThreadPool(int required_thread_num, bool forced = false);

~ThreadPool();

// Rgister invoker in thread pool

int RegisterInvoker(const std::string& worker_name = {});

int UnregisterInvoker(int invoker_id);

int PostTask(int invoker_id, FuncType&& func);

// Note: this function spends extra time on creating packaged tasks.

// We should avoid using this function for posting small and repeated tasks.

// Use it only when you want synchronous invocation.

template <typename F, typename... Args>

auto PostTaskWithRes(int invoker_id, F&& f, Args&&... args)

-> std::future<typename std::result_of<F(Args...)>::type> {

using task_ret_type = typename std::result_of<F(Args...)>::type;

auto package = std::make_shared<std::packaged_task<task_ret_type()>>(

std::bind(std::forward<F>(f), std::forward<Args>(args)...));

auto task = [package] { (*package)();};

std::thread::id th_id;

{

std::lock_guard<std::mutex> thread_lock(thread_mutex_);

auto err = findThreadId(invoker_id, th_id);

if (err != 0) {

return {}; // causing ret future in an invalid state

}

}

auto res = package->get_future();

std::lock_guard<std::mutex> task_lock(task_mutex_);

if (insertTask(invoker_id, th_id, std::move(task)) != 0) {

return {}; // causing ret future in an invalid state

}

return res;

}

private:

ThreadPool(const ThreadPool&) = delete;

ThreadPool(ThreadPool&&) = delete;

ThreadPool& operator=(const ThreadPool&) = delete;

ThreadPool& operator=(ThreadPool&&) = delete;

int generatorInvokerId();

void initThread(int invoker_id, const std::string& th_name);

std::thread::id findLeastBusyThread();

int findThreadId(int invoker_id, std::thread::id& th_id);

int insertTask(int invoker_id, std::thread::id th_id, FuncType&& task);

private:

struct Task {

int invoker_id = 0;

FuncType task = nullptr;

};

std::mutex thread_mutex_;

std::vector<std::thread> threads_;

// thread name to (id, reference_count) map

std::unordered_map<std::string, std::pair<std::thread::id, int>> thread_name_id_map_;

// invoker_id to thread_name map

std::unordered_map<int, std::string> registered_invokers_;

std::mutex task_mutex_;

// thread_id to task list map

std::unordered_map<std::thread::id, std::list<Task>> assigned_tasks_;

std::list<Task> unassigned_tasks_;

std::condition_variable task_notifier_;

int max_thread_num_ = {0};

bool stop_all_ = false;

};

} // namespace extensions

} // namespace agora