Line data    Source code

       1              : #ifndef THINGER_ASIO_WORKERS
       2              : #define THINGER_ASIO_WORKERS
       3              : 
       4              : #include <unordered_map>
       5              : #include <set>
       6              : #include <boost/asio.hpp>
       7              : #include "worker_thread.hpp"
       8              : #include "worker_client.hpp"
       9              : 
      10              : namespace thinger::asio {
      11              : 
      12              :         class workers {
      13              :         public:
      14              :         workers();
      15              :                 virtual ~workers();
      16              : 
      17              :                 /// start the asio workers
      18              :         bool start(size_t working_threads=std::thread::hardware_concurrency());
      19              : 
      20              :         /// stop the asio workers and all their pending async operations
      21              :         bool stop();
      22              : 
      23              :         /// keep the asio workers alive until a signal is received
      24              :         void wait(const std::set<unsigned>& signals = {SIGINT, SIGTERM, SIGQUIT});
      25              :         
      26              :         /// check if workers are running
      27           30 :         bool running() const { return running_; }
      28              : 
      29              :         /// return an isolated io_context not shared in the pool
      30              :         boost::asio::io_context& get_isolated_io_context(std::string thread_name);
      31              : 
      32              :         /// return the next io_context (round-robin)
      33              :         boost::asio::io_context& get_next_io_context();
      34              : 
      35              :         /// return the io_context associated with the caller thread
      36              :                 boost::asio::io_context& get_thread_io_context();
      37              : 
      38              :         // Client management
      39              :         /// Register a client that uses workers
      40              :         void register_client(worker_client* client);
      41              :         
      42              :         /// Unregister a client
      43              :         void unregister_client(worker_client* client);
      44              :         
      45              :         /// Get number of registered clients
      46              :         size_t client_count() const;
      47              :         
      48              :         /// Enable/disable automatic management based on clients
      49              :         void set_auto_manage(bool enable) { auto_manage_ = enable; }
      50              :         
      51              :         /// Check if auto management is enabled
      52            9 :         bool is_auto_managed() const { return auto_manage_; }
      53              : 
      54              :         private:
      55              :         /// Internal method to perform the actual stop
      56              :         void do_stop();
      57              :         
      58              :         /// mutex used for initializing threads and data structures
      59              :         std::mutex mutex_;
      60              : 
      61              :         // io_context used for capturing signals and wait coordination
      62              :         boost::asio::io_context wait_context_;
      63              :         boost::asio::signal_set signals_;
      64              :         
      65              :         // Work guard to keep wait_context_ running
      66              :         using work_guard_type = boost::asio::executor_work_guard<boost::asio::io_context::executor_type>;
      67              :         std::unique_ptr<work_guard_type> wait_work_;
      68              : 
      69              :         /// flag for controlling if it is running
      70              :                 std::atomic<bool> running_{false};
      71              : 
      72              :         /// index to control next io_context to use (in the round-robin)
      73              :         unsigned next_io_context_   = 0;
      74              : 
      75              :                 /// worker threads used for general asio pool
      76              :                 std::vector<std::unique_ptr<worker_thread>> worker_threads_;
      77              : 
      78              :         /// worker threads allocated from isolated_io_contexts
      79              :         std::vector<std::unique_ptr<worker_thread>> job_threads_;
      80              : 
      81              :                 /// relation of all worker threads with their worker thread instance
      82              :         std::unordered_map<std::thread::id, std::reference_wrapper<worker_thread>> workers_threads_map_;
      83              : 
      84              :         // Client management
      85              :         /// Set of registered clients
      86              :         std::set<worker_client*> clients_;
      87              :         
      88              :         /// Mutex for client management
      89              :         mutable std::mutex clients_mutex_;
      90              :         
      91              :         /// Flag to enable/disable automatic management
      92              :         std::atomic<bool> auto_manage_{true};
      93              : 
      94              :         };
      95              : 
      96              :         // Singleton instance accessor
      97              :         workers& get_workers();
      98              : 
      99              : }
     100              : 
     101              : #endif