Add ability to specify CPU affinity for wrk threads

Makefile

@@ -14,7 +14,7 @@ else ifeq ($(TARGET), darwin)
 	LIBS += -L/usr/local/opt/openssl/lib
 	CFLAGS += -I/usr/local/include -I/usr/local/opt/openssl/include
 else ifeq ($(TARGET), linux)
-        CFLAGS  += -D_POSIX_C_SOURCE=200809L -D_BSD_SOURCE
+        CFLAGS  += -D_GNU_SOURCE
 	LIBS    += -ldl
 	LDFLAGS += -Wl,-E
 else ifeq ($(TARGET), freebsd)

src/units.c

@@ -4,9 +4,12 @@
 #include <stdio.h>
 #include <strings.h>
 #include <inttypes.h>
+#include <errno.h>
+#include <string.h>
 
 #include "units.h"
 #include "aprintf.h"
+#include "zmalloc.h"
 
 typedef struct {
     int scale;
@@ -102,3 +105,40 @@ int scan_metric(char *s, uint64_t *n) {
 int scan_time(char *s, uint64_t *n) {
     return scan_units(s, n, &time_units_s);
 }
+
+int scan_affinity(char *s,  struct aff_set_head **as_head) {
+    char *str = strdup(s);
+    char *token = NULL;
+    char *saveptr;
+    struct aff_set_head *head = NULL;
+
+    token = strtok_r(str, ",", &saveptr);
+    if (token != NULL) {
+        do {
+            int val;
+            struct aff_set *item = NULL;
+
+            errno = 0;
+            val = (int)strtol(token, NULL, 0);
+            if (errno != 0) {
+                free(str);
+                return -1;
+            }
+
+            item = zcalloc(sizeof(struct aff_set));
+            CPU_ZERO(&item->set);
+            CPU_SET(val, &item->set);
+            if (head == NULL) {
+                head = zmalloc(sizeof(struct aff_set_head));
+                STAILQ_INIT(head);
+            }
+            STAILQ_INSERT_TAIL(head, item, items);
+
+            token = strtok_r(NULL, ",", &saveptr);
+        } while (token != NULL);
+    }
+
+    free(str);
+    *as_head = head;
+    return 0;
+}

src/units.h

@@ -1,6 +1,9 @@
 #ifndef UNITS_H
 #define UNITS_H
 
+#include <sys/queue.h>
+#include <sched.h>
+
 char *format_binary(long double);
 char *format_metric(long double);
 char *format_time_us(long double);
@@ -9,4 +12,13 @@ char *format_time_s(long double);
 int scan_metric(char *, uint64_t *);
 int scan_time(char *, uint64_t *);
 
+struct aff_set {
+    cpu_set_t set;
+    STAILQ_ENTRY(aff_set) items;
+};
+
+STAILQ_HEAD(aff_set_head, aff_set);
+
+int scan_affinity(char *, struct aff_set_head **);
+
 #endif /* UNITS_H */

src/wrk.c

@@ -5,12 +5,15 @@
 #include "main.h"
 #include "hdr_histogram.h"
 #include "stats.h"
+#include "units.h"
+#include <assert.h>
 
 // Max recordable latency of 1 day
 #define MAX_LATENCY 24L * 60 * 60 * 1000000
 
 static struct config {
     uint64_t threads;
+    struct aff_set_head *affinity;
     uint64_t connections;
     uint64_t duration;
     uint64_t timeout;
@@ -55,6 +58,8 @@ static void usage() {
            "    -c, --connections <N>  Connections to keep open   \n"
            "    -d, --duration    <T>  Duration of test           \n"
            "    -t, --threads     <N>  Number of threads to use   \n"
+           "    -a, --affinity    <S>  Threads affinity: comma    \n"
+           "                           separated list of CPUs     \n"
            "                                                      \n"
            "    -s, --script      <S>  Load Lua script file       \n"
            "    -H, --header      <H>  Add header to request      \n"
@@ -123,13 +128,23 @@ int main(int argc, char **argv) {
     uint64_t connections = cfg.connections / cfg.threads;
     double throughput    = (double)cfg.rate / cfg.threads;
     uint64_t stop_at     = time_us() + (cfg.duration * 1000000);
+    struct aff_set_head *as_head = cfg.affinity;
+    struct aff_set      *as_item;
+    if (as_head != NULL)
+       as_item = STAILQ_FIRST(as_head);
 
     for (uint64_t i = 0; i < cfg.threads; i++) {
         thread *t = &threads[i];
         t->loop        = aeCreateEventLoop(10 + cfg.connections * 3);
         t->connections = connections;
         t->throughput = throughput;
         t->stop_at     = stop_at;
+        t->cpu_set = NULL;
+        if (as_head != NULL) {
+            assert(as_item != NULL);
+            t->cpu_set = &as_item->set;
+            as_item = STAILQ_NEXT(as_item, items);
+        }
 
         t->L = script_create(cfg.script, url, headers);
         script_init(L, t, argc - optind, &argv[optind]);
@@ -251,6 +266,17 @@ void *thread_main(void *arg) {
     thread *thread = arg;
     aeEventLoop *loop = thread->loop;
 
+    if (thread->cpu_set != NULL) {
+        int res;
+        cpu_set_t *set = thread->cpu_set;
+
+        res = pthread_setaffinity_np(pthread_self(), sizeof(*set), set);
+        if (res != 0) {
+            fprintf(stderr, "set thread affinity failed (errno: %d)\n", res);
+            exit(EXIT_FAILURE);
+        }
+    }
+
     thread->cs = zcalloc(thread->connections * sizeof(connection));
     tinymt64_init(&thread->rand, time_us());
     hdr_init(1, MAX_LATENCY, 3, &thread->latency_histogram);
@@ -695,6 +721,7 @@ static struct option longopts[] = {
     { "connections",    required_argument, NULL, 'c' },
     { "duration",       required_argument, NULL, 'd' },
     { "threads",        required_argument, NULL, 't' },
+    { "affinity",       required_argument, NULL, 'a' },
     { "script",         required_argument, NULL, 's' },
     { "header",         required_argument, NULL, 'H' },
     { "latency",        no_argument,       NULL, 'L' },
@@ -712,17 +739,21 @@ static int parse_args(struct config *cfg, char **url, struct http_parser_url *pa
 
     memset(cfg, 0, sizeof(struct config));
     cfg->threads     = 2;
+    cfg->affinity    = NULL;
     cfg->connections = 10;
     cfg->duration    = 10;
     cfg->timeout     = SOCKET_TIMEOUT_MS;
     cfg->rate        = 0;
     cfg->record_all_responses = true;
 
-    while ((c = getopt_long(argc, argv, "t:c:d:s:H:T:R:LUBrv?", longopts, NULL)) != -1) {
+    while ((c = getopt_long(argc, argv, "t:a:c:d:s:H:T:R:LUBrv?", longopts, NULL)) != -1) {
         switch (c) {
             case 't':
                 if (scan_metric(optarg, &cfg->threads)) return -1;
                 break;
+            case 'a':
+                if (scan_affinity(optarg, &cfg->affinity)) return -1;
+                break;
             case 'c':
                 if (scan_metric(optarg, &cfg->connections)) return -1;
                 break;
@@ -782,6 +813,20 @@ static int parse_args(struct config *cfg, char **url, struct http_parser_url *pa
         return -1;
     }
 
+    if (cfg->affinity != NULL) {
+        uint64_t i = 0;
+        struct aff_set *set;
+
+        STAILQ_FOREACH(set, cfg->affinity, items)
+            i++;
+
+        if (cfg->threads != i) {
+            fprintf(stderr, "number of affinity CPUs (%lu) does not match the "
+                    "number of threads (%lu)\n", i, cfg->threads);
+            return -1;
+        }
+    }
+
     *url    = argv[optind];
     *header = NULL;
 

src/wrk.h

@@ -30,6 +30,7 @@ typedef struct {
     aeEventLoop *loop;
     struct addrinfo *addr;
     uint64_t connections;
+    cpu_set_t *cpu_set;
     int interval;
     uint64_t stop_at;
     uint64_t complete;