From 04a25b1bfa7fabdfce63af6a1124434bccd957d4 Mon Sep 17 00:00:00 2001 From: Damir Mansurov Date: Mon, 22 Aug 2022 15:42:56 +0300 Subject: [PATCH] Add ability to specify CPU affinity for wrk threads In some cases, it is necessary that the wrk2 threads run on a specific CPU. Example of usage with affinity: "wrk -t 3 -a 0,1,2 -c 100 -d 5 -R 500000 ". If the affinity parameter is not specified, then the threads will be run as before (without any affinitization). --- Makefile | 2 +- src/units.c | 40 ++++++++++++++++++++++++++++++++++++++++ src/units.h | 12 ++++++++++++ src/wrk.c | 47 ++++++++++++++++++++++++++++++++++++++++++++++- src/wrk.h | 1 + 5 files changed, 100 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index a537a68..cb0df17 100644 --- a/Makefile +++ b/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) diff --git a/src/units.c b/src/units.c index f94c6a7..88f879b 100644 --- a/src/units.c +++ b/src/units.c @@ -4,9 +4,12 @@ #include #include #include +#include +#include #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; +} diff --git a/src/units.h b/src/units.h index 625f9ad..77ed1f6 100644 --- a/src/units.h +++ b/src/units.h @@ -1,6 +1,9 @@ #ifndef UNITS_H #define UNITS_H +#include +#include + 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 */ diff --git a/src/wrk.c b/src/wrk.c index 1049f0b..c3484d7 100644 --- a/src/wrk.c +++ b/src/wrk.c @@ -5,12 +5,15 @@ #include "main.h" #include "hdr_histogram.h" #include "stats.h" +#include "units.h" +#include // 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 Connections to keep open \n" " -d, --duration Duration of test \n" " -t, --threads Number of threads to use \n" + " -a, --affinity Threads affinity: comma \n" + " separated list of CPUs \n" " \n" " -s, --script Load Lua script file \n" " -H, --header Add header to request \n" @@ -123,6 +128,10 @@ 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]; @@ -130,6 +139,12 @@ int main(int argc, char **argv) { 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; diff --git a/src/wrk.h b/src/wrk.h index 202a4c3..aab35ad 100644 --- a/src/wrk.h +++ b/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;