diff --git a/CMakeLists.txt b/CMakeLists.txt index d1edf00..1d1dc74 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -52,6 +52,7 @@ if(WIN32) target_link_options(ncmdump_lib PRIVATE -static) endif() else() + set(CMAKE_CXX_FLAGS -pthread) add_executable(ncmdump_exec ${COMMON_HEADERS} ${COMMON_SOURCES} diff --git a/src/include/thread_pool.hpp b/src/include/thread_pool.hpp new file mode 100644 index 0000000..90193f9 --- /dev/null +++ b/src/include/thread_pool.hpp @@ -0,0 +1,73 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace cxxpool +{ + +class ThreadPool +{ +public: + ThreadPool(): + thread_num_(std::thread::hardware_concurrency() * 2), + is_running_(true) + { + work_threads_ = std::make_unique(thread_num_); + for (unsigned int i = 0; i < thread_num_; ++i) + { + work_threads_[i] = std::thread([&](){ + while (true) + { + std::unique_lock lock(mutex_); + cv_.wait(lock, [&](){ return !is_running_ || !queue_.empty(); }); + if (!is_running_ && queue_.empty()) + return; + else if (queue_.empty()) + continue; + + std::function func = std::move(queue_.front()); + queue_.pop(); + lock.unlock(); + std::invoke(func); + } + }); + } + } + + ~ThreadPool() noexcept + { + is_running_ = false; + cv_.notify_all(); + + for (unsigned int i = 0; i < thread_num_; ++i) + { + if (work_threads_[i].joinable()) + work_threads_[i].join(); + } + } + + template + void submit(Func&& func, Args&&... args) + { + std::unique_lock lock(mutex_); + queue_.push(std::bind(func, std::forward(args)...)); + lock.unlock(); + cv_.notify_one(); + } + +private: + const unsigned int thread_num_; + std::unique_ptr work_threads_; + std::mutex mutex_; + std::atomic is_running_; + std::condition_variable cv_; + std::queue> queue_; +}; + +} // namespace cxxpool diff --git a/src/main.cpp b/src/main.cpp index 1c1cf0b..feceb9d 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -12,39 +12,46 @@ #include "color.h" #include "version.h" #include "cxxopts.hpp" +#include "thread_pool.hpp" namespace fs = std::filesystem; +cxxpool::ThreadPool thread_pool; + void processFile(const fs::path &filePath, const fs::path &outputFolder) { - if (fs::exists(filePath) == false) - { - std::cerr << BOLDRED << "[Error] " << RESET << "file '" << filePath.u8string() << "' does not exist." << std::endl; - return; - } + auto work_func = [](fs::path filePath, fs::path outputFolder) { + if (fs::exists(filePath) == false) + { + std::cerr << BOLDRED << "[Error] " << RESET << "file '" << filePath.u8string() << "' does not exist." << std::endl; + return; + } - // skip if not ending with ".ncm" - if (!filePath.has_extension() || filePath.extension().u8string() != ".ncm") - { - return; - } + // skip if not ending with ".ncm" + if (!filePath.has_extension() || filePath.extension().u8string() != ".ncm") + { + return; + } - try - { - NeteaseCrypt crypt(filePath.u8string()); - crypt.Dump(outputFolder.u8string()); - crypt.FixMetadata(); + try + { + NeteaseCrypt crypt(filePath.u8string()); + crypt.Dump(outputFolder.u8string()); + crypt.FixMetadata(); - std::cout << BOLDGREEN << "[Done] " << RESET << "'" << filePath.u8string() << "' -> '" << crypt.dumpFilepath().u8string() << "'" << std::endl; - } - catch (const std::invalid_argument &e) - { - std::cerr << BOLDRED << "[Exception] " << RESET << RED << e.what() << RESET << " '" << filePath.u8string() << "'" << std::endl; - } - catch (...) - { - std::cerr << BOLDRED << "[Error] Unexpected exception while processing file: " << RESET << filePath.u8string() << std::endl; - } + std::cout << BOLDGREEN << "[Done] " << RESET << "'" << filePath.u8string() << "' -> '" << crypt.dumpFilepath().u8string() << "'" << std::endl; + } + catch (const std::invalid_argument &e) + { + std::cerr << BOLDRED << "[Exception] " << RESET << RED << e.what() << RESET << " '" << filePath.u8string() << "'" << std::endl; + } + catch (...) + { + std::cerr << BOLDRED << "[Error] Unexpected exception while processing file: " << RESET << filePath.u8string() << std::endl; + } + }; + + thread_pool.submit(work_func, filePath, outputFolder); } int main(int argc, char **argv)