Auto-detect number of threads

Wiebe Cazemier
1 parent 64ec702a
Showing 3 changed files with 28 additions and 4 deletions
mainapp.cpp
mainapp.h
threaddata.cpp
@@ -4,6 +4,7 @@
 #include "getopt.h"
 #include <unistd.h>
 #include <stdio.h>
+#include <sys/sysinfo.h>
  
 #include <openssl/ssl.h>
 #include <openssl/err.h>
@@ -11,7 +12,6 @@
 #include "logger.h"
  
 #define MAX_EVENTS 1024
-#define NR_OF_THREADS 4
  
 #define VERSION "0.1"
  
@@ -148,12 +148,20 @@ void do_thread_work(ThreadData *threadData)
 MainApp::MainApp(const std::string &configFilePath) :
     subscriptionStore(new SubscriptionStore())
 {
+    this->num_threads = get_nprocs();
+
+    if (num_threads <= 0)
+        throw std::runtime_error("Invalid number of CPUs: " + std::to_string(num_threads));
+
     epollFdAccept = check<std::runtime_error>(epoll_create(999));
     taskEventFd = eventfd(0, EFD_NONBLOCK);
  
     confFileParser.reset(new ConfigFileParser(configFilePath));
     loadConfig();
  
+    // TODO: override in conf possibility.
+    logger->logf(LOG_NOTICE, "%d CPUs are detected, making as many threads.", num_threads);
+
     auto f = std::bind(&MainApp::queueCleanup, this);
     timer.addCallback(f, 86400000, "session expiration");
 }
@@ -342,7 +350,7 @@ void MainApp::start()
     ev.events = EPOLLIN;
     check<std::runtime_error>(epoll_ctl(this->epollFdAccept, EPOLL_CTL_ADD, taskEventFd, &ev));
  
-    for (int i = 0; i < NR_OF_THREADS; i++)
+    for (int i = 0; i < num_threads; i++)
     {
         std::shared_ptr<ThreadData> t(new ThreadData(i, subscriptionStore, *confFileParser.get()));
         t->start(&do_thread_work);
@@ -373,7 +381,7 @@ void MainApp::start()
             {
                 if (cur_fd == listen_fd_plain || cur_fd == listen_fd_ssl)
                 {
-                    std::shared_ptr<ThreadData> thread_data = threads[next_thread_index++ % NR_OF_THREADS];
+                    std::shared_ptr<ThreadData> thread_data = threads[next_thread_index++ % num_threads];
  
                     logger->logf(LOG_INFO, "Accepting connection on thread %d", thread_data->threadnr);
  
@@ -24,6 +24,7 @@
 class MainApp
 {
     static MainApp *instance;
+    int num_threads = 0;
  
     bool started = false;
     bool running = true;
@@ -21,7 +21,22 @@ void ThreadData::start(thread_f f)
     std::ostringstream threadName;
     threadName << "FlashMQ T " << threadnr;
     threadName.flush();
-    pthread_setname_np(native, threadName.str().c_str());
+    const char *c_str = threadName.str().c_str();
+    pthread_setname_np(native, c_str);
+
+    cpu_set_t cpuset;
+    CPU_ZERO(&cpuset);
+    CPU_SET(threadnr, &cpuset);
+    check<std::runtime_error>(pthread_setaffinity_np(native, sizeof(cpuset), &cpuset));
+
+    // It's not really necessary to get affinity again, but now I'm logging truth instead assumption.
+    check<std::runtime_error>(pthread_getaffinity_np(native, sizeof(cpuset), &cpuset));
+    int pinned_cpu = -1;
+    for (int j = 0; j < CPU_SETSIZE; j++)
+        if (CPU_ISSET(j, &cpuset))
+            pinned_cpu = j;
+
+    logger->logf(LOG_NOTICE, "Thread '%s' pinned to CPU %d", c_str, pinned_cpu);
 }
  
 void ThreadData::quit()