diff --git a/src/agent/CMakeLists.txt b/src/agent/CMakeLists.txt index 90816ff477..1b0aa0cb77 100644 --- a/src/agent/CMakeLists.txt +++ b/src/agent/CMakeLists.txt @@ -12,15 +12,21 @@ include_directories(include) set(SOURCES src/agent.cpp + src/task_manager.cpp ) set(HEADERS include/agent.hpp include/itask_manager.hpp + include/task_manager.hpp ) add_library(agent ${SOURCES} ${HEADERS}) +find_package(Boost REQUIRED COMPONENTS asio) + +target_link_libraries(agent PRIVATE ${Boost_LIBRARIES}) + if(BUILD_TESTS) add_subdirectory(tests) endif() diff --git a/src/agent/include/task_manager.hpp b/src/agent/include/task_manager.hpp new file mode 100644 index 0000000000..303c049116 --- /dev/null +++ b/src/agent/include/task_manager.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include + +#include +#include + +#include +#include +#include + +class TaskManager : public ITaskManager> +{ +public: + TaskManager(); + + void start(size_t numThreads) override; + void stop() override; + + void enqueueTask(std::function task) override; + void enqueueTask(boost::asio::awaitable task) override; + +private: + boost::asio::io_context m_ioContext; + boost::asio::io_context::work m_work; + std::vector m_threads; +}; diff --git a/src/agent/src/task_manager.cpp b/src/agent/src/task_manager.cpp new file mode 100644 index 0000000000..7378cac158 --- /dev/null +++ b/src/agent/src/task_manager.cpp @@ -0,0 +1,45 @@ +#include + +#include +#include +#include + +TaskManager::TaskManager() + : m_work(m_ioContext) +{ +} + +void TaskManager::start(size_t numThreads) +{ + stop(); + + for (size_t i = 0; i < numThreads; ++i) + { + m_threads.emplace_back([this]() { m_ioContext.run(); }); + } +} + +void TaskManager::stop() +{ + m_ioContext.stop(); + + for (std::thread& thread : m_threads) + { + if (thread.joinable()) + { + thread.join(); + } + } + m_threads.clear(); + m_ioContext.reset(); +} + +void TaskManager::enqueueTask(std::function task) +{ + boost::asio::post(m_ioContext, std::move(task)); +} + +void TaskManager::enqueueTask(boost::asio::awaitable task) +{ + boost::asio::co_spawn(m_ioContext, std::move(task), boost::asio::detached); +} diff --git a/src/agent/tests/CMakeLists.txt b/src/agent/tests/CMakeLists.txt index e06cc1d3bb..c9a9fcfb64 100644 --- a/src/agent/tests/CMakeLists.txt +++ b/src/agent/tests/CMakeLists.txt @@ -4,8 +4,13 @@ find_package(GTest CONFIG REQUIRED) set(TEST_SOURCES agent_test.cpp + task_manager_test.cpp ) add_executable(agent_test agent_test.cpp) target_link_libraries(agent_test PRIVATE agent GTest::gtest) add_test(NAME AgentTest COMMAND agent_test) + +add_executable(task_manager_test task_manager_test.cpp) +target_link_libraries(task_manager_test PRIVATE agent GTest::gtest) +add_test(NAME TaskManagerTest COMMAND task_manager_test) diff --git a/src/agent/tests/task_manager_test.cpp b/src/agent/tests/task_manager_test.cpp new file mode 100644 index 0000000000..3cb74cc43e --- /dev/null +++ b/src/agent/tests/task_manager_test.cpp @@ -0,0 +1,66 @@ +#include +#include + +#include +#include + +class TaskManagerTest : public ::testing::Test +{ +protected: + TaskManager taskManager; + + void SetUp() override {} + + void TearDown() override {} +}; + +TEST_F(TaskManagerTest, StartAndStop) +{ + taskManager.start(2); + taskManager.stop(); +} + +TEST_F(TaskManagerTest, EnqueueFunctionTask) +{ + taskManager.start(1); + + std::atomic counter = 0; + std::function task = [&counter]() + { + ++counter; + }; + + taskManager.enqueueTask(task); + + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + + EXPECT_EQ(counter, 1); + + taskManager.stop(); +} + +TEST_F(TaskManagerTest, EnqueueCoroutineTask) +{ + taskManager.start(1); + + std::atomic counter = 0; + auto coroutineTask = [&counter]() -> boost::asio::awaitable + { + ++counter; + co_return; + }; + + taskManager.enqueueTask(coroutineTask()); + + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + + EXPECT_EQ(counter, 1); + + taskManager.stop(); +} + +int main(int argc, char** argv) +{ + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/vcpkg.json b/src/vcpkg.json index 3b45724a08..aeda316c92 100644 --- a/src/vcpkg.json +++ b/src/vcpkg.json @@ -2,6 +2,7 @@ "name": "wazuh-agent-mvp", "version": "5.0.0", "dependencies": [ + "boost-asio", "gtest" ] }