numam/util/memloadgen.cc

#include <sys/endian.h>
#include <sys/select.h>
#include <sys/signal.h>
#include "gen.hh"
#include <array>
#include <atomic>
#include <cstdlib>
#include <cstring>
#include <list>
#include <iostream>
#include <fstream>
#include "ntr.h"
#include "nms.h"
#include <getopt.h>
#include <pthread.h>
#include <unistd.h>
#include <topo.h>

static void
usage()
{
	fprintf(stdout,
	    "Usage:\n"
		"    -v: verbose mode\n"
	    "    -b: buffer size\n"
		"    -q: bytes per second\n"
	    "    -d: destination domain index\n"
	    "    -s: worker threads cpu list\n"
		"    -S: enable shared buffer\n"
		"    -t: time to run\n"
		"    -T: transaction size\n"
		"    -i: inter arrival time distribution\n"
		"    -o: output file path\n"
		"    -H: history size for pct adjustment\n"
		"    -M: print this string when threads are ready to run\n");
	fflush(stdout);
}

static char output_file[256] = "memloadgen_samples.txt";

int main(int argc, char * argv[])
{
	ntr_init();
	ntr_set_level(NTR_DEP_USER1, NTR_LEVEL_WARNING);

	size_t arr_sz = 64 * 1024 * 1024;
	uint32_t time = -1;
	uint64_t bps = 0;
	uint64_t transaction_size = arr_sz;
	cpuset_t threads;
	char magic[256] = {0};
	CPU_ZERO(&threads);
	CPU_SET(0, &threads);
	char ia_dist[32] = "fixed";
	int history_sz = 5;
	std::list<uint64_t> history;

	int shared_buffer = 0;
	int rate_ctrl = 0;
	cpuset_t domain_mask;
	CPU_ZERO(&domain_mask);
	CPU_SET(0, &domain_mask);
	{
		int c;
		// parse arguments
		while ((c = getopt(argc, argv, "vhb:d:s:So:T:t:q:i:H:M:")) != -1) {
			switch (c) {
			case 'v':
				ntr_set_level(NTR_DEP_USER1, ntr_get_level(NTR_DEP_USER1) + 1);
				break;
			case 'h':
				usage();
				exit(0);
			case 'b':
				arr_sz = strtoull(optarg, nullptr, 10);
				break;
			case 'd':
				cpulist_to_cpuset(optarg, &domain_mask);
				break;
			case 's':
				cpulist_to_cpuset(optarg, &threads);
				break;
			case 'S':
				shared_buffer = 1;
				break;
			case 'o':
				strncpy(output_file, optarg, 256);
				break;
			case 't':
				time = strtoul(optarg, nullptr, 10);
				break;
			case 'T':
				transaction_size = strtoul(optarg, nullptr, 10);
				break;
			case 'q':
				bps = (uint64_t)strtoull(optarg, nullptr, 10);
				break;
			case 'i':
				strncpy(ia_dist, optarg, sizeof(ia_dist));
				break;
			case 'H':
				history_sz = strtol(optarg, nullptr, 10);
				break;
			case 'M':
				strncpy(magic, optarg, sizeof(magic));
				break;
			default:
				usage();
				exit(0);
			}
		}
	}

	ntr(NTR_DEP_USER1, NTR_LEVEL_INFO, "Configruation:\n"
									   "    buffer size: %ld\n"
									   "    num threads: %d\n"
									   "    target domain: %ld\n"
									   "    bytes per second: %lu\n"
									   "    interarrival distribution: %s\n"
									   "    shared buffer: %d\n"
									   "    transaction time: %lu\n"
									   "    runtime: %d\n"
									   "    history: %d\n"
									   "    magic: %s\n",
									    arr_sz, CPU_COUNT(&threads), 
										CPU_FFS(&domain_mask) - 1, bps, 
										ia_dist, shared_buffer, 
										transaction_size,time, history_sz, magic);

	// init topo
	if (topo_init(ntr_get_level(NTR_DEP_USER1) != NTR_LEVEL_DEFAULT)) {
		fprintf(stderr, "libtopo init failed!\n");
		exit(1);
	}

	// init 
	if (nms_init(ntr_get_level(NTR_DEP_USER1) != NTR_LEVEL_DEFAULT)) {
		fprintf(stderr, "libnms init failed!\n");
		exit(1);
	}

	bool success = false;
	memload_generator::memload_generator_options opts;
	opts.buffer_size = arr_sz;
	opts.trans_per_second = bps / transaction_size;
	opts.shared_buffer = shared_buffer;
	opts.transaction_size = transaction_size;
	opts.verbose = ntr_get_level(NTR_DEP_USER1) != NTR_LEVEL_DEFAULT;
	strncpy(opts.ia_dist, ia_dist, sizeof(opts.ia_dist));
	std::ofstream ofile;
	ofile.open(output_file, std::ios::out | std::ios::trunc);

    auto mgen = new memload_generator(&threads, &domain_mask, &opts, &success);
	if (strlen(magic) > 0) {
		fprintf(stdout, "%s\n", magic);
		fflush(stdout);
	}
	if (!mgen->start()) {
		fprintf(stderr, "failed to start memloadgen!\n");
		exit(1);
	}

	struct timeval stval;
	stval.tv_sec = 0;
	stval.tv_usec = 0;
	char pct_line[64] = {0};

	uint64_t prev_ts = topo_uptime_ns();
	uint64_t prev_trans = mgen->get_transactions();
	uint32_t cur_time = 0;
	while(cur_time < time) {
		usleep(S2US);
		uint64_t cur_ts = topo_uptime_ns();
		uint64_t trans = mgen->get_transactions();
		uint64_t bps = (uint64_t)((double)((trans - prev_trans) * transaction_size) / ((double)(cur_ts - prev_ts) / (double)S2NS));
			
		ntr(NTR_DEP_USER1, NTR_LEVEL_INFO, "%ldB,%ldM\n", bps, bps / 1024 / 1024);
		ofile << "s," << cur_time << "," << bps << std::endl;
		ofile.flush();
		
		prev_ts = cur_ts;
		prev_trans = trans;
		cur_time++;

		if (rate_ctrl == 0) {
			// keep history
			history.emplace_back(bps);
			if ((int)history.size() > history_sz) {
				history.pop_front();
			}

			fd_set fdset;
			FD_ZERO(&fdset);
			FD_SET(STDIN_FILENO, &fdset);
			int ret = select(1, &fdset, NULL, NULL, &stval);
			if (ret < 0) {
				if (errno != EINTR) {
					fprintf(stderr, "select() failed with %d\n", errno);
					exit(1);
				}
			} else if (ret > 0) {
				if (FD_ISSET(STDIN_FILENO, &fdset)) {
					ret = read(STDIN_FILENO, pct_line, sizeof(pct_line) - 1);
					if (ret < 0) {
						fprintf(stderr, "read() failed with %d\n", errno);
						exit(1);
					}
					unsigned int pct = strtoul(pct_line, NULL, 10);
					uint64_t sum = 0;
					size_t sz = history.size();
					while (history.size() > 0) {
						sum += history.front();
						history.pop_front();
					}

					uint64_t newbps = ((sum / sz) * (double)pct / 100.0);
					mgen->set_transactions(newbps / transaction_size);
					ntr(NTR_DEP_USER1, NTR_LEVEL_INFO, "adjusted target bps to %u%% = %ldB ~= %ldM\n", pct, newbps, newbps / 1024 / 1024);
					
					ofile << "p," << cur_time << "," << pct << std::endl;
					ofile.flush();

					rate_ctrl = 1;
				}
			}
		}
	}
	mgen->stop();
	delete mgen;
	ofile.close();

	return 0;
}