Implement -A flag to set CPU affinity.

src/iperf.h

@@ -144,6 +144,7 @@ struct iperf_test
     int       omit;                             /* duration of omit period (-O flag) */
     int       duration;                         /* total duration of test (-t flag) */
     char     *diskfile_name;			/* -F option */
+    int       affinity, server_affinity;	/* -A option */
 
     int       ctrl_sck;
     int       listener;

src/iperf3.1

@@ -34,6 +34,12 @@ of using random data;
 server-side: read from the network and write to the file, instead
 of throwing the data away
 .TP
+.BR -A ", " --affinity " \fIn/n,m\fR"
+Set the CPU affinity, if possible (linux only).
+On both the client and server you can set the local affinity;
+in addition, on the client side you can override the server's
+affinity for just that one test, using the n,m form.
+.TP
 .BR -V ", " --verbose " "
 give more detailed output 
 .TP

src/iperf_api.c

@@ -7,6 +7,9 @@
  * for complete information.
  */
 
+#define _GNU_SOURCE
+#define __USE_GNU
+
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -520,6 +523,7 @@ iperf_parse_arguments(struct iperf_test *test, int argc, char **argv)
         {"zerocopy", no_argument, NULL, 'Z'},
         {"omit", required_argument, NULL, 'O'},
         {"file", required_argument, NULL, 'F'},
+        {"affinity", required_argument, NULL, 'A'},
         {"help", no_argument, NULL, 'h'},
 
     /*  XXX: The following ifdef needs to be split up. linux-congestion is not
@@ -533,10 +537,11 @@ iperf_parse_arguments(struct iperf_test *test, int argc, char **argv)
     int flag;
     int blksize;
     int server_flag, client_flag, rate_flag;
+    char* comma;
 
     blksize = 0;
     server_flag = client_flag = rate_flag = 0;
-    while ((flag = getopt_long(argc, argv, "p:f:i:DVJdvsc:ub:t:n:l:P:Rw:B:M:N46S:L:ZO:F:h", longopts, NULL)) != -1) {
+    while ((flag = getopt_long(argc, argv, "p:f:i:DVJdvsc:ub:t:n:l:P:Rw:B:M:N46S:L:ZO:F:A:h", longopts, NULL)) != -1) {
         switch (flag) {
             case 'p':
                 test->server_port = atoi(optarg);
@@ -687,6 +692,22 @@ iperf_parse_arguments(struct iperf_test *test, int argc, char **argv)
             case 'F':
                 test->diskfile_name = optarg;
                 break;
+            case 'A':
+                test->affinity = atoi(optarg);
+                if (test->affinity < 0 || test->affinity > 1024) {
+                    i_errno = IEAFFINITY;
+                    return -1;
+                }
+		comma = strchr(optarg, ',');
+		if (comma != NULL) {
+		    test->server_affinity = atoi(comma+1);
+		    if (test->server_affinity < 0 || test->server_affinity > 1024) {
+			i_errno = IEAFFINITY;
+			return -1;
+		    }
+		    client_flag = 1;
+		}
+                break;
             case 'h':
             default:
                 usage_long();
@@ -988,6 +1009,8 @@ send_parameters(struct iperf_test *test)
 	else if (test->protocol->id == Pudp)
 	    cJSON_AddTrueToObject(j, "udp");
 	cJSON_AddIntToObject(j, "omit", test->omit);
+	if (test->server_affinity != -1)
+	    cJSON_AddIntToObject(j, "server_affinity", test->server_affinity);
 	if (test->duration)
 	    cJSON_AddIntToObject(j, "time", test->duration);
 	if (test->settings->bytes)
@@ -1038,6 +1061,8 @@ get_parameters(struct iperf_test *test)
 	    set_protocol(test, Pudp);
 	if ((j_p = cJSON_GetObjectItem(j, "omit")) != NULL)
 	    test->omit = j_p->valueint;
+	if ((j_p = cJSON_GetObjectItem(j, "server_affinity")) != NULL)
+	    test->server_affinity = j_p->valueint;
 	if ((j_p = cJSON_GetObjectItem(j, "time")) != NULL)
 	    test->duration = j_p->valueint;
 	if ((j_p = cJSON_GetObjectItem(j, "num")) != NULL)
@@ -1351,6 +1376,8 @@ iperf_defaults(struct iperf_test *testp)
     testp->omit = OMIT;
     testp->duration = DURATION;
     testp->diskfile_name = (char*) 0;
+    testp->affinity = -1;
+    testp->server_affinity = -1;
     testp->server_port = PORT;
     testp->ctrl_sck = -1;
     testp->prot_listener = -1;
@@ -1495,7 +1522,7 @@ iperf_reset_test(struct iperf_test *test)
     set_protocol(test, Ptcp);
     test->omit = OMIT;
     test->duration = DURATION;
-    test->diskfile_name = (char*) 0;
+    test->server_affinity = -1;
     test->state = 0;
     test->server_hostname = NULL;
 
@@ -2142,3 +2169,47 @@ iperf_json_finish(struct iperf_test *test)
     test->json_top = test->json_start = test->json_intervals = test->json_end = NULL;
     return 0;
 }
+
+
+/* CPU affinity stuff - linux only. */
+
+int
+iperf_setaffinity(int affinity)
+{
+#ifdef linux
+    cpu_set_t cpu_set;
+
+    CPU_ZERO(&cpu_set);
+    CPU_SET(affinity, &cpu_set);
+    if (sched_setaffinity(0, sizeof(cpu_set_t), &cpu_set) != 0) {
+	i_errno = IEAFFINITY;
+        return -1;
+    }
+    return 0;
+#else /*linux*/
+    i_errno = IEAFFINITY;
+    return -1;
+#endif /*linux*/
+}
+
+int
+iperf_clearaffinity(void)
+{
+#ifdef linux
+    cpu_set_t cpu_set;
+    int i;
+
+    CPU_ZERO(&cpu_set);
+    for (i = 0; i < CPU_SETSIZE; ++i)
+	CPU_SET(i, &cpu_set);
+    if (sched_setaffinity(0, sizeof(cpu_set_t), &cpu_set) != 0) {
+	i_errno = IEAFFINITY;
+        return -1;
+    }
+    return 0;
+#else /*linux*/
+    i_errno = IEAFFINITY;
+    return -1;
+#endif /*linux*/
+}
+

src/iperf_api.h

@@ -209,6 +209,10 @@ void iperf_test_reset(struct iperf_test *);
 int iperf_json_start(struct iperf_test *);
 int iperf_json_finish(struct iperf_test *);
 
+/* CPU affinity routines */
+int iperf_setaffinity(int affinity);
+int iperf_clearaffinity(void);
+
 /* Error routines. */
 void iperf_err(struct iperf_test *test, const char *format, ...) __attribute__ ((format(printf,2,3)));
 void iperf_errexit(struct iperf_test *test, const char *format, ...) __attribute__ ((format(printf,2,3),noreturn));
@@ -264,6 +268,7 @@ enum {
     IENONBLOCKING = 129,    // Unable to set socket to non-blocking (check perror)
     IESETWINDOWSIZE = 130,  // Unable to set socket window size (check perror)
     IEPROTOCOL = 131,       // Protocol does not exist
+    IEAFFINITY = 132,       // Unable to set CPU affinity (check perror)
     /* Stream errors */
     IECREATESTREAM = 200,   // Unable to create a new stream (check herror/perror)
     IEINITSTREAM = 201,     // Unable to initialize stream (check herror/perror)

src/iperf_client_api.c

@@ -313,6 +313,10 @@ iperf_run_client(struct iperf_test * test)
     fd_set read_set, write_set;
     struct timeval now;
 
+    if (test->affinity != -1)
+	if (iperf_setaffinity(test->affinity) != 0)
+	    return -1;
+
     if (test->json_output)
 	if (iperf_json_start(test) < 0)
 	    return -1;
@@ -326,9 +330,8 @@ iperf_run_client(struct iperf_test * test)
     }
 
     /* Start the client and connect to the server */
-    if (iperf_connect(test) < 0) {
+    if (iperf_connect(test) < 0)
         return -1;
-    }
 
     /* Begin calculating CPU utilization */
     cpu_util(NULL);

src/iperf_error.c

@@ -231,6 +231,10 @@ iperf_strerror(int i_errno)
         case IEPROTOCOL:
             snprintf(errstr, len, "protocol does not exist");
             break;
+        case IEAFFINITY:
+            snprintf(errstr, len, "unable to set CPU affinity");
+            perr = 1;
+            break;
         case IECREATESTREAM:
             snprintf(errstr, len, "unable to create a new stream");
             herr = 1;

src/iperf_server_api.c

@@ -137,6 +137,9 @@ iperf_accept(struct iperf_test *test)
             return -1;
         if (iperf_exchange_parameters(test) < 0)
             return -1;
+	if (test->server_affinity != -1) 
+	    if (iperf_setaffinity(test->server_affinity) != 0)
+		return -1;
         if (test->on_connect)
             test->on_connect(test);
     } else {
@@ -255,6 +258,8 @@ iperf_test_reset(struct iperf_test *test)
     test->omit = OMIT;
     test->duration = DURATION;
     test->diskfile_name = (char*) 0;
+    test->affinity = -1;
+    test->server_affinity = -1;
     test->state = 0;
     test->server_hostname = NULL;
 
@@ -335,6 +340,10 @@ iperf_run_server(struct iperf_test *test)
     struct iperf_stream *sp;
     struct timeval now;
 
+    if (test->affinity != -1) 
+	if (iperf_setaffinity(test->affinity) != 0)
+	    return -1;
+
     if (test->json_output)
 	if (iperf_json_start(test) < 0)
 	    return -1;
@@ -486,5 +495,9 @@ iperf_run_server(struct iperf_test *test)
 	    return -1;
     } 
 
+    if (test->server_affinity != -1) 
+	if (iperf_clearaffinity() != 0)
+	    return -1;
+
     return 0;
 }

src/locale.c

@@ -72,6 +72,8 @@ const char usage_longstr[] = "Usage: iperf [-s|-c host] [options]\n"
                            "  -p, --port      #         server port to listen on/connect to\n"
                            "  -f, --format    [kmgKMG]  format to report: Kbits, Mbits, KBytes, MBytes\n"
                            "  -i, --interval  #         seconds between periodic bandwidth reports\n"
+                           "  -F, --file name           xmit/recv the specified file\n"
+                           "  -A, --affinity n/n,m      set CPU affinity\n"
                            "  -V, --verbose             more detailed output\n"
                            "  -J, --json                output in JSON format\n"
                            "  -v, --version             show version information and quit\n"
@@ -99,7 +101,6 @@ const char usage_longstr[] = "Usage: iperf [-s|-c host] [options]\n"
                            /* "  -L, --flowlabel N         set the IPv6 'flow label'\n" */
                            "  -Z, --zerocopy            use a 'zero copy' method of sending data\n"
                            "  -O, --omit N              omit the first n seconds\n"
-                           "  -F, --file name           xmit/recv the specified file\n"
 
 #ifdef NOT_YET_SUPPORTED /* still working on these */
                            "  -D, --daemon              run the server as a daemon\n"