From 9ece0fe0209673d2f3f7f4f3e3738dfa7ec1873b Mon Sep 17 00:00:00 2001
From: Vadim Kochan <vadim4j@gmail.com>
Date: Wed, 23 Dec 2015 22:31:23 +0200
Subject: trafgen: Add option to specify packets sending rate

Added -b,--rate option in units of:

    pps/B/kB/MB/kBit/Mbit/Gbit/KiB/MiB/GiB

to specify rate at which packets will be sent.
Similarly to -t,--gap option the packets will be sent
in slow mode with 1 CPU.

Tested with ifpps.

Signed-off-by: Vadim Kochan <vadim4j@gmail.com>
Signed-off-by: Tobias Klauser <tklauser@distanz.ch>
---
 trafgen.8 |   4 ++
 trafgen.c | 140 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 135 insertions(+), 9 deletions(-)

diff --git a/trafgen.8 b/trafgen.8
index 6f5d318..57ce8fa 100644
--- a/trafgen.8
+++ b/trafgen.8
@@ -136,6 +136,10 @@ probing for a given interval, iii) ping-like debugging with specific payload
 patterns. Furthermore, the TX_RING interface does not cope with interpacket
 gaps.
 .PP
+.SS -b <rate>, --rate <rate>
+Specify the packet send rate <num>pps/B/kB/MB/GB/kBit/Mbit/Gbit/KiB/MiB/GiB units.
+Like with the -t,--gap option, the packets are sent in slow mode.
+.PP
 .SS -S <size>, --ring-size <size>
 Manually define the TX_RING resp. TX_RING size in ''<num>KiB/MiB/GiB''. On
 default the size is being determined based on the network connectivity rate.
diff --git a/trafgen.c b/trafgen.c
index df2ecf4..a1d5463 100644
--- a/trafgen.c
+++ b/trafgen.c
@@ -55,6 +55,27 @@
 #include "ring_tx.h"
 #include "csum.h"
 
+#ifndef timeval_to_timespec
+#define timeval_to_timespec(tv, ts) {         \
+	(ts)->tv_sec = (tv)->tv_sec;          \
+	(ts)->tv_nsec = (tv)->tv_usec * 1000; \
+}
+#endif
+
+enum shaper_type {
+	SHAPER_NONE,
+	SHAPER_PKTS,
+	SHAPER_BYTES,
+};
+
+struct shaper {
+	enum shaper_type type;
+	unsigned long long sent;
+	unsigned long long rate;
+	struct timeval start;
+	struct timeval end;
+};
+
 struct ctx {
 	bool rand, rfraw, jumbo_support, verbose, smoke_test, enforce, qdisc_path;
 	size_t reserve_size;
@@ -64,6 +85,7 @@ struct ctx {
 	char *device, *device_trans, *rhost;
 	struct timespec gap;
 	struct sockaddr_in dest;
+	struct shaper sh;
 	char *packet_str;
 };
 
@@ -83,7 +105,7 @@ size_t plen = 0;
 struct packet_dyn *packet_dyn = NULL;
 size_t dlen = 0;
 
-static const char *short_options = "d:c:n:t:vJhS:rk:i:o:VRs:P:eE:pu:g:CHQqD:";
+static const char *short_options = "d:c:n:t:vJhS:rk:i:o:VRs:P:eE:pu:g:CHQqD:b:";
 static const struct option long_options[] = {
 	{"dev",			required_argument,	NULL, 'd'},
 	{"out",			required_argument,	NULL, 'o'},
@@ -91,6 +113,7 @@ static const struct option long_options[] = {
 	{"conf",		required_argument,	NULL, 'c'},
 	{"num",			required_argument,	NULL, 'n'},
 	{"gap",			required_argument,	NULL, 't'},
+	{"rate",		required_argument,	NULL, 'b'},
 	{"cpus",		required_argument,	NULL, 'P'},
 	{"ring-size",		required_argument,	NULL, 'S'},
 	{"kernel-pull",		required_argument,	NULL, 'k'},
@@ -172,6 +195,7 @@ static void __noreturn help(void)
 	     "  -r|--rand                      Randomize packet selection (def: round robin)\n"
 	     "  -P|--cpus <uint>               Specify number of forks(<= CPUs) (def: #CPUs)\n"
 	     "  -t|--gap <time>                Set approx. interpacket gap (s/ms/us/ns, def: us)\n"
+	     "  -b|--rate <rate>               Send traffic at specified rate (pps/B/kB/MB/GB/kBit/Mbit/Gbit/KiB/MiB/GiB)\n"
 	     "  -S|--ring-size <size>          Manually set mmap size (KiB/MiB/GiB)\n"
 	     "  -E|--seed <uint>               Manually set srand(3) seed\n"
 	     "  -u|--user <userid>             Drop privileges and change to userid\n"
@@ -535,6 +559,52 @@ static int xmit_smoke_probe(int icmp_sock, struct ctx *ctx)
 	return -1;
 }
 
+static bool shaper_is_set(struct shaper *sh)
+{
+	return sh->type != SHAPER_NONE;
+}
+
+static void shaper_start(struct shaper *sh)
+{
+	bug_on(gettimeofday(&sh->start, NULL));
+	sh->sent = 0;
+}
+
+static void shaper_init(struct shaper *sh, unsigned long long rate, enum shaper_type type)
+{
+	memset(sh, 0, sizeof(struct shaper));
+	sh->rate = rate;
+	sh->type = type;
+}
+
+static void shaper_delay(struct shaper *sh, unsigned long pkt_len)
+{
+	if ((sh->start.tv_sec | sh->start.tv_usec) <= 0)
+		return;
+
+	sh->sent += sh->type == SHAPER_BYTES ? pkt_len : 1;
+
+	if (sh->sent >= sh->rate) {
+		struct timeval delay_us;
+		struct timespec delay_ns;
+		struct timeval time_sent;
+		struct timeval time_1s = { .tv_sec = 1 };
+
+		bug_on(gettimeofday(&sh->end, NULL));
+		timersub(&sh->end, &sh->start, &time_sent);
+
+		if (timercmp(&time_1s, &time_sent, > )) {
+			timersub(&time_1s, &time_sent, &delay_us);
+			timeval_to_timespec(&delay_us, &delay_ns);
+
+			if ((delay_ns.tv_sec | delay_ns.tv_nsec) > 0)
+				nanosleep(&delay_ns, NULL);
+		}
+
+		shaper_start(sh);
+	}
+}
+
 static void xmit_slowpath_or_die(struct ctx *ctx, unsigned int cpu, unsigned long orig_num)
 {
 	int ret, icmp_sock = -1;
@@ -560,6 +630,9 @@ static void xmit_slowpath_or_die(struct ctx *ctx, unsigned int cpu, unsigned lon
 
 	bug_on(gettimeofday(&start, NULL));
 
+	if (shaper_is_set(&ctx->sh))
+		shaper_start(&ctx->sh);
+
 	while (likely(sigint == 0 && num > 0 && plen > 0)) {
 		pktd = &packet_dyn[i];
 		if (packet_dyn_has_elems(pktd)) {
@@ -605,6 +678,9 @@ retry:
 		if (ctx->num > 0)
 			num--;
 
+		if (shaper_is_set(&ctx->sh))
+			shaper_delay(&ctx->sh, packets[i].len);
+
 		if ((ctx->gap.tv_sec | ctx->gap.tv_nsec) > 0)
 			nanosleep(&ctx->gap, NULL);
 	}
@@ -901,6 +977,8 @@ int main(int argc, char **argv)
 	int min_opts = 5;
 	char **cpp_argv = NULL;
 	size_t cpp_argc = 0;
+	unsigned long long rate;
+	enum shaper_type shape_type;
 
 	fmemset(&ctx, 0, sizeof(ctx));
 	ctx.cpus = get_number_cpus_online();
@@ -999,9 +1077,7 @@ int main(int argc, char **argv)
 			ctx.num = orig_num;
 			break;
 		case 't':
-			slow = true;
 			ptr = optarg;
-			prctl(PR_SET_TIMERSLACK, 1UL);
 			gap = strtoul(optarg, NULL, 0);
 
 			for (j = i = strlen(optarg); i > 0; --i) {
@@ -1025,14 +1101,51 @@ int main(int argc, char **argv)
 			} else if (!strncmp(ptr, "s", strlen("s"))) {
 				ctx.gap.tv_sec = gap;
 				ctx.gap.tv_nsec = 0;
-			} else
+			} else {
 				panic("Syntax error in time param!\n");
+			}
+			break;
+		case 'b':
+			rate = strtoul(optarg, &ptr, 0);
+			if (!rate || optarg == ptr)
+				panic("Invalid rate param\n");
+
+			if (strncmp(ptr, "pps", strlen("pps")) == 0) {
+				shape_type = SHAPER_PKTS;
+			} else if (strncmp(ptr, "B", strlen("B")) == 0) {
+				shape_type = SHAPER_BYTES;
+			} else if (strncmp(ptr, "kB", strlen("kB")) == 0) {
+				shape_type = SHAPER_BYTES;
+				rate *= 1000;
+			} else if (strncmp(ptr, "MB", strlen("MB")) == 0) {
+				shape_type = SHAPER_BYTES;
+				rate *= 1000 * 1000;
+			} else if (strncmp(ptr, "GB", strlen("GB")) == 0) {
+				shape_type = SHAPER_BYTES;
+				rate *= 1000 * 1000 * 1000;
+			} else if (strncmp(ptr, "kbit", strlen("kbit")) == 0) {
+				shape_type = SHAPER_BYTES;
+				rate *= 1000 / 8;
+			} else if (strncmp(ptr, "Mbit", strlen("Mbit")) == 0) {
+				shape_type = SHAPER_BYTES;
+				rate *= 1000 * 1000 / 8;
+			} else if (strncmp(ptr, "Gbit", strlen("Gbit")) == 0) {
+				shape_type = SHAPER_BYTES;
+				rate *= 1000 * 1000 * 1000 / 8;
+			} else if (strncmp(ptr, "KiB", strlen("KiB")) == 0) {
+				shape_type = SHAPER_BYTES;
+				rate *= 1 << 10;
+			} else if (strncmp(ptr, "MiB", strlen("MiB")) == 0) {
+				shape_type = SHAPER_BYTES;
+				rate *= 1 << 20;
+			} else if (strncmp(ptr, "GiB", strlen("GiB")) == 0) {
+				shape_type = SHAPER_BYTES;
+				rate *= 1 << 30;
+			} else {
+				panic("Invalid unit type for rate\n");
+			}
 
-			if (gap > 0)
-				/* Fall back to single core to not mess up
-				 * correct timing. We are slow anyway!
-				 */
-				ctx.cpus = 1;
+			shaper_init(&ctx.sh, rate, shape_type);
 			break;
 		case 'S':
 			ptr = optarg;
@@ -1118,6 +1231,15 @@ int main(int argc, char **argv)
 		sleep(0);
 	}
 
+	if (shaper_is_set(&ctx.sh) || (ctx.gap.tv_sec | ctx.gap.tv_nsec) > 0) {
+		prctl(PR_SET_TIMERSLACK, 1UL);
+		/* Fall back to single core to not mess up correct timing.
+		 * We are slow anyway!
+		 */
+		ctx.cpus = 1;
+		slow = true;
+	}
+
 	/*
 	 * If number of packets is smaller than number of CPUs use only as
 	 * many CPUs as there are packets. Otherwise we end up sending more
-- 
cgit v1.2.3-54-g00ecf