dtoctl.c

/* Copyright (C) 2026 SK hynix, Inc. */
/* SPDX-License-Identifier: MIT */

/*
 * dtoctl -- a CLI wrapper that runs a program with the DSA Transparent Offload
 * (DTO) library preloaded.  It converts command line options into the DTO_*
 * environment variables and prepends libdto.so to LD_PRELOAD so that users do
 * not need to export the variables or set up LD_PRELOAD manually.
 */

#ifndef _GNU_SOURCE
#define _GNU_SOURCE     /* for asprintf() and environ */
#endif

#include <argp.h>
#include <ctype.h>
#include <errno.h>
#include <libgen.h>
#include <limits.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#ifndef LIBDTO_PATH
#define LIBDTO_PATH "/usr/lib64/libdto.so"
#endif

#ifndef DTO_VERSION
#define DTO_VERSION "unknown"
#endif

/* exit statuses for dtoctl's own failures, following env(1) and timeout(1) */
#define EXIT_DTOCTL_FAILURE	125
#define EXIT_CANNOT_INVOKE	126
#define EXIT_ENOENT		127

/* libdto copies DTO_WQ_LIST into a buffer of this size */
#define WQ_LIST_MAX		256

/* argp option keys for the long-only options (no short equivalents) */
enum {
	OPT_UMWAIT_DELAY = 0x100,
	OPT_OVERLAPPING_MEMMOVE,
	OPT_LOG_FILE,
	OPT_LOG_LEVEL,
	OPT_STATS_FILE,
	OPT_NO_MEMCPY,
	OPT_NO_MEMMOVE,
	OPT_NO_MEMSET,
	OPT_NO_MEMCMP,
	OPT_NO_CACHE_CONTROL,
	OPT_NO_AUTO_ADJUST,
	OPT_STDC_ONLY,
	OPT_STATS,
	OPT_DRY_RUN,
};

struct opts {
	int idx;
	char *exename;

	/* library / preload */
	const char *library;
	bool ignore_env;
	bool dry_run;

	/* value options (NULL means "not specified") */
	const char *wait_method;
	const char *min_bytes;
	const char *cpu_fraction;
	const char *numa_aware;
	const char *wq_list;
	const char *umwait_delay;
	const char *overlapping_memmove; /* "0" or "1" after validation */
	const char *log_file;
	const char *log_level;
	const char *stats_file;

	/* flag options (false means "not specified") */
	bool no_memcpy;
	bool no_memmove;
	bool no_memset;
	bool no_memcmp;
	bool no_cache_control;
	bool no_auto_adjust;
	bool stdc_only;
	bool stats;
};

/* (a part of) output in --help and --version options (generated by argp runtime) */
const char *argp_program_version = "dtoctl (DTO) " DTO_VERSION;
const char *argp_program_bug_address = "https://github.com/intel/DTO/issues";

static const char *const wait_methods[] = {
	"busypoll", "umwait", "yield", "tpause", NULL
};

/* Option groups are separated using entries with only a .doc field. */
static struct argp_option dtoctl_options[] = {
	{.doc = "General:", .group = 1},
	{.name = "library",
	 .key = 'l',
	 .arg = "PATH",
	 .doc = "Path to libdto.so to prepend to LD_PRELOAD. Overrides $DTO_LIBRARY "
		"and the installed library (" LIBDTO_PATH ")"},
	{.name = "ignore-env",
	 .key = 'i',
	 .doc = "Unset the DTO_* variables inherited from the environment, so that "
		"only the options given here and DTO's defaults apply"},
	{.name = "dry-run",
	 .key = OPT_DRY_RUN,
	 .doc = "Print the library, the DTO_* environment and the command that "
		"would run, then exit without running it"},

	{.doc = "Offload tuning:", .group = 2},
	{.name = "wait-method",
	 .key = 'w',
	 .arg = "METHOD",
	 .doc = "How to wait for DSA completion: yield, busypoll, umwait, or tpause "
		"(sets DTO_WAIT_METHOD; default: busypoll)"},
	{.name = "min-bytes",
	 .key = 'b',
	 .arg = "BYTES",
	 .doc = "Minimum operation size offloaded to DSA; accepts a K, M or G suffix "
		"(sets DTO_MIN_BYTES; default: 64K)"},
	{.name = "cpu-fraction",
	 .key = 'c',
	 .arg = "FRACTION",
	 .doc = "Fraction (0.0 <= f < 1.0) of each operation done on CPU in parallel "
		"to DSA (sets DTO_CPU_SIZE_FRACTION; default: 0.0)"},
	{.name = "numa-aware",
	 .key = 'n',
	 .arg = "MODE",
	 .doc = "NUMA awareness: 0 disable, 1 buffer-centric, 2 cpu-centric "
		"(sets DTO_IS_NUMA_AWARE; default: 0)"},
	{.name = "wq-list",
	 .key = 'q',
	 .arg = "LIST",
	 .doc = "Semicolon-separated list of DSA WQs to use, e.g. \"wq0.0;wq2.0\" "
		"(sets DTO_WQ_LIST; default: auto-discover all WQs)"},
	{.name = "umwait-delay",
	 .key = OPT_UMWAIT_DELAY,
	 .arg = "CYCLES",
	 .doc = "Delay in cycles for the umwait instruction (sets DTO_UMWAIT_DELAY; "
		"default: 100000)"},
	{.name = "overlapping-memmove",
	 .key = OPT_OVERLAPPING_MEMMOVE,
	 .arg = "WHERE",
	 .doc = "Where to run memmove with overlapping buffers: cpu or dsa "
		"(sets DTO_OVERLAPPING_MEMMOVE_ACTION; default: cpu)"},

	{.doc = "Disable DSA offload for specific operations (DSA is used by default):",
	 .group = 3},
	{.name = "no-memcpy",
	 .key = OPT_NO_MEMCPY,
	 .doc = "Use system memcpy instead of DSA (sets DTO_DSA_MEMCPY=0)"},
	{.name = "no-memmove",
	 .key = OPT_NO_MEMMOVE,
	 .doc = "Use system memmove instead of DSA (sets DTO_DSA_MEMMOVE=0)"},
	{.name = "no-memset",
	 .key = OPT_NO_MEMSET,
	 .doc = "Use system memset instead of DSA (sets DTO_DSA_MEMSET=0)"},
	{.name = "no-memcmp",
	 .key = OPT_NO_MEMCMP,
	 .doc = "Use system memcmp instead of DSA (sets DTO_DSA_MEMCMP=0)"},

	{.doc = "Other toggles:", .group = 4},
	{.name = "no-cache-control",
	 .key = OPT_NO_CACHE_CONTROL,
	 .doc = "Clear the DSA cache control flag to avoid cache pollution "
		"(sets DTO_DSA_CC=0; default: cache control on)"},
	{.name = "no-auto-adjust",
	 .key = OPT_NO_AUTO_ADJUST,
	 .doc = "Disable auto tuning of cpu-fraction and min-bytes "
		"(sets DTO_AUTO_ADJUST_KNOBS=0; default: auto tuning on)"},
	{.name = "stdc-only",
	 .key = OPT_STDC_ONLY,
	 .doc = "Use only the standard C memory functions, no DSA offload "
		"(sets DTO_USESTDC_CALLS=1; default: DSA offload enabled)"},

	{.doc = "Statistics and logging:", .group = 5},
	{.name = "stats",
	 .key = OPT_STATS,
	 .doc = "Enable stats collection; for debugging/profiling only "
		"(sets DTO_COLLECT_STATS=1; default: off)"},
	{.name = "stats-file",
	 .key = OPT_STATS_FILE,
	 .arg = "PATH",
	 .doc = "Write the stats histogram to PATH instead of standard output; "
		"implies --stats (sets DTO_STATS_FILE; default: standard output)"},
	{.name = "log-file",
	 .key = OPT_LOG_FILE,
	 .arg = "PATH",
	 .doc = "Redirect DTO output to PATH.<program>.<pid> (sets DTO_LOG_FILE; "
		"default: standard output)"},
	{.name = "log-level",
	 .key = OPT_LOG_LEVEL,
	 .arg = "LEVEL",
	 .doc = "Verbosity of DTO logging: 0, 1, or 2 (sets DTO_LOG_LEVEL; "
		"default: 0)"},

	{0},
};

/* Parse a decimal integer, rejecting signs, empty strings and trailing garbage. */
static bool parse_ulong(const char *s, unsigned long *val)
{
	char *end;

	if (!isdigit((unsigned char)*s))
		return false;
	errno = 0;
	*val = strtoul(s, &end, 10);
	return errno == 0 && *end == '\0';
}

/* Parse a byte count with an optional K, M or G suffix (powers of 1024). */
static bool parse_size(const char *s, unsigned long *val)
{
	unsigned int shift = 0;
	unsigned long n;
	char *end;

	if (!isdigit((unsigned char)*s))
		return false;
	errno = 0;
	n = strtoul(s, &end, 10);
	if (errno)
		return false;

	switch (*end) {
	case 'k': case 'K':
		shift = 10;
		end++;
		break;
	case 'm': case 'M':
		shift = 20;
		end++;
		break;
	case 'g': case 'G':
		shift = 30;
		end++;
		break;
	}
	if (*end != '\0' || n > (ULONG_MAX >> shift))
		return false;

	*val = n << shift;
	return true;
}

/* Accept an integer in [0, max] for options that select a mode or level. */
static const char *check_level(struct argp_state *state, const char *name,
			       char *arg, unsigned long max)
{
	unsigned long val;

	if (!parse_ulong(arg, &val) || val > max)
		argp_error(state, "invalid %s '%s': must be an integer from 0 to %lu",
			   name, arg, max);
	return arg;
}

static const char *check_path(struct argp_state *state, const char *name, char *arg)
{
	if (*arg == '\0')
		argp_error(state, "--%s needs a non-empty PATH", name);
	return arg;
}

static error_t parse_option(int key, char *arg, struct argp_state *state)
{
	struct opts *opts = state->input;
	unsigned long val = 0;
	char *end, *bytes;
	double frac;
	int i;

	switch (key) {
	case 'l':
		opts->library = check_path(state, "library", arg);
		break;
	case 'i':
		opts->ignore_env = true;
		break;
	case OPT_DRY_RUN:
		opts->dry_run = true;
		break;
	case 'w':
		for (i = 0; wait_methods[i]; i++) {
			if (!strcmp(arg, wait_methods[i]))
				break;
		}
		if (wait_methods[i] == NULL)
			argp_error(state, "invalid wait method '%s': must be yield, "
				   "busypoll, umwait, or tpause", arg);
		opts->wait_method = arg;
		break;
	case 'b':
		if (!parse_size(arg, &val))
			argp_error(state, "invalid minimum size '%s': must be a byte "
				   "count, optionally with a K, M or G suffix", arg);
		/* libdto takes plain bytes, so expand any suffix here */
		if (asprintf(&bytes, "%lu", val) < 0)
			argp_failure(state, EXIT_DTOCTL_FAILURE, errno, "--min-bytes");
		opts->min_bytes = bytes;
		break;
	case 'c':
		errno = 0;
		frac = strtod(arg, &end);
		/* the negated range check also rejects NaN */
		if (end == arg || *end != '\0' || errno || !(frac >= 0.0 && frac < 1.0))
			argp_error(state, "invalid cpu fraction '%s': must be a number "
				   ">= 0.0 and < 1.0", arg);
		opts->cpu_fraction = arg;
		break;
	case 'n':
		opts->numa_aware = check_level(state, "NUMA mode", arg, 2);
		break;
	case 'q':
		if (*arg == '\0' || strlen(arg) >= WQ_LIST_MAX)
			argp_error(state, "invalid WQ list: must be 1 to %d characters",
				   WQ_LIST_MAX - 1);
		opts->wq_list = arg;
		break;
	case OPT_UMWAIT_DELAY:
		/* libdto treats 0 as "use the default", so reject it */
		if (!parse_ulong(arg, &val) || val == 0)
			argp_error(state, "invalid umwait delay '%s': must be a positive "
				   "number of cycles", arg);
		opts->umwait_delay = arg;
		break;
	case OPT_OVERLAPPING_MEMMOVE:
		if (!strcmp(arg, "cpu"))
			opts->overlapping_memmove = "0";
		else if (!strcmp(arg, "dsa"))
			opts->overlapping_memmove = "1";
		else
			argp_error(state, "invalid --overlapping-memmove '%s': must be "
				   "cpu or dsa", arg);
		break;
	case OPT_NO_MEMCPY:
		opts->no_memcpy = true;
		break;
	case OPT_NO_MEMMOVE:
		opts->no_memmove = true;
		break;
	case OPT_NO_MEMSET:
		opts->no_memset = true;
		break;
	case OPT_NO_MEMCMP:
		opts->no_memcmp = true;
		break;
	case OPT_NO_CACHE_CONTROL:
		opts->no_cache_control = true;
		break;
	case OPT_NO_AUTO_ADJUST:
		opts->no_auto_adjust = true;
		break;
	case OPT_STDC_ONLY:
		opts->stdc_only = true;
		break;
	case OPT_STATS:
		opts->stats = true;
		break;
	case OPT_LOG_FILE:
		opts->log_file = check_path(state, "log-file", arg);
		break;
	case OPT_LOG_LEVEL:
		opts->log_level = check_level(state, "log level", arg, 2);
		break;
	case OPT_STATS_FILE:
		opts->stats_file = check_path(state, "stats-file", arg);
		break;

	case ARGP_KEY_ARG:
		/* leave the program and its arguments to ARGP_KEY_ARGS */
		return ARGP_ERR_UNKNOWN;

	case ARGP_KEY_ARGS:
		/* process remaining non-option arguments */
		opts->exename = state->argv[state->next];
		opts->idx = state->next;
		break;

	case ARGP_KEY_END:
		/* --dry-run can show the environment without a program */
		if (opts->exename == NULL && !opts->dry_run)
			argp_usage(state);
		break;

	default:
		return ARGP_ERR_UNKNOWN;
	}
	return 0;
}

/* DTO_* variables set from options, in the order they were set, for --dry-run */
static const char *option_vars[32];
static int nr_option_vars;

static void set_option_var(const char *name, const char *value)
{
	if (setenv(name, value, 1)) {
		perror("dtoctl: setenv");
		exit(EXIT_DTOCTL_FAILURE);
	}
	option_vars[nr_option_vars++] = name;
}

static bool is_option_var(const char *name, int len)
{
	for (int i = 0; i < nr_option_vars; i++) {
		if ((int)strlen(option_vars[i]) == len && !strncmp(option_vars[i], name, len))
			return true;
	}
	return false;
}

/* Unset every DTO_* variable inherited from the caller (--ignore-env). */
static void clear_dto_environ(void)
{
	char **env, *eq, *name;
	bool found;

	/* unsetenv() reshuffles environ, so rescan it after every removal */
	do {
		found = false;
		for (env = environ; *env; env++) {
			eq = strchr(*env, '=');
			if (strncmp(*env, "DTO_", 4) || eq == NULL)
				continue;
			name = strndup(*env, eq - *env);
			if (name == NULL || unsetenv(name)) {
				perror("dtoctl: unsetenv");
				exit(EXIT_DTOCTL_FAILURE);
			}
			free(name);
			found = true;
			break;
		}
	} while (found);
}

static void setup_dto_environ(struct opts *opts)
{
	/* value options: only set when explicitly specified */
	if (opts->wait_method)
		set_option_var("DTO_WAIT_METHOD", opts->wait_method);
	if (opts->min_bytes)
		set_option_var("DTO_MIN_BYTES", opts->min_bytes);
	if (opts->cpu_fraction)
		set_option_var("DTO_CPU_SIZE_FRACTION", opts->cpu_fraction);
	if (opts->numa_aware)
		set_option_var("DTO_IS_NUMA_AWARE", opts->numa_aware);
	if (opts->wq_list)
		set_option_var("DTO_WQ_LIST", opts->wq_list);
	if (opts->umwait_delay)
		set_option_var("DTO_UMWAIT_DELAY", opts->umwait_delay);
	if (opts->overlapping_memmove)
		set_option_var("DTO_OVERLAPPING_MEMMOVE_ACTION", opts->overlapping_memmove);
	if (opts->log_file)
		set_option_var("DTO_LOG_FILE", opts->log_file);
	if (opts->log_level)
		set_option_var("DTO_LOG_LEVEL", opts->log_level);
	if (opts->stats_file)
		set_option_var("DTO_STATS_FILE", opts->stats_file);

	/* flag options: only set when explicitly specified */
	if (opts->no_memcpy)
		set_option_var("DTO_DSA_MEMCPY", "0");
	if (opts->no_memmove)
		set_option_var("DTO_DSA_MEMMOVE", "0");
	if (opts->no_memset)
		set_option_var("DTO_DSA_MEMSET", "0");
	if (opts->no_memcmp)
		set_option_var("DTO_DSA_MEMCMP", "0");
	if (opts->no_cache_control)
		set_option_var("DTO_DSA_CC", "0");
	if (opts->no_auto_adjust)
		set_option_var("DTO_AUTO_ADJUST_KNOBS", "0");
	if (opts->stdc_only)
		set_option_var("DTO_USESTDC_CALLS", "1");
	/* libdto ignores DTO_STATS_FILE unless stats are collected */
	if (opts->stats || opts->stats_file)
		set_option_var("DTO_COLLECT_STATS", "1");
}

/* The libdto.so installed relative to this dtoctl, if there is one. */
static char *installed_library(void)
{
#ifdef LIBDTO_REL_PATH
	char exe[PATH_MAX], *path, *real;
	ssize_t len;

	len = readlink("/proc/self/exe", exe, sizeof(exe) - 1);
	if (len < 0)
		return NULL;
	exe[len] = '\0';

	if (asprintf(&path, "%s/%s", dirname(exe), LIBDTO_REL_PATH) < 0)
		return NULL;
	real = realpath(path, NULL);
	free(path);
	if (real && access(real, R_OK)) {
		free(real);
		real = NULL;
	}
	return real;
#else
	return NULL;
#endif
}

/*
 * Resolve the libdto.so to preload: --library, then $DTO_LIBRARY, then the
 * library installed alongside dtoctl, then the compiled-in default.  The
 * result is an absolute path, so it still loads after the program changes
 * directory.  A missing library is an error here; ld.so would only print a
 * warning and run the program without DTO.
 */
static char *resolve_library(struct opts *opts, const char **source)
{
	const char *lib;
	char *path;

	if (opts->library) {
		lib = opts->library;
		*source = "--library";
	} else if ((lib = getenv("DTO_LIBRARY")) && *lib) {
		*source = "DTO_LIBRARY";
	} else if ((path = installed_library())) {
		*source = "installed with dtoctl";
		return path;
	} else {
		lib = LIBDTO_PATH;
		*source = "compiled-in default";
	}

	path = realpath(lib, NULL);
	if (path == NULL || access(path, R_OK)) {
		fprintf(stderr, "dtoctl: cannot use DTO library '%s' (%s): %s\n",
			lib, *source, strerror(errno));
		if (opts->library == NULL)
			fprintf(stderr, "dtoctl: use -l/--library or DTO_LIBRARY to "
				"choose a libdto.so\n");
		exit(EXIT_DTOCTL_FAILURE);
	}

	/* ld.so splits LD_PRELOAD at colons and spaces */
	if (strpbrk(path, ": ")) {
		fprintf(stderr, "dtoctl: DTO library path '%s' cannot be preloaded: "
			"it contains a colon or a space\n", path);
		exit(EXIT_DTOCTL_FAILURE);
	}
	return path;
}

/* true for an LD_PRELOAD entry that names a libdto.so (libdto.so, libdto.so.1, ...) */
static bool is_libdto(const char *entry)
{
	const char *base = strrchr(entry, '/');

	base = base ? base + 1 : entry;
	return !strncmp(base, "libdto.so", strlen("libdto.so"));
}

/*
 * Put lib at the front of LD_PRELOAD.  Other libdto.so entries are dropped:
 * two copies of DTO in one process would both initialize and both open the
 * work queues.
 */
static void setup_preload(const char *lib)
{
	const char *cur = getenv("LD_PRELOAD");
	char *preload, *copy, *entry, *save, *real;

	/*
	 * LD_PRELOAD is a list, so it can be much longer than a single path.
	 * Size the result from the inputs: the kept entries and their colons
	 * never take more than the old value plus one byte.
	 */
	preload = malloc(strlen(lib) + (cur ? strlen(cur) : 0) + 2);
	copy = strdup(cur ? cur : "");
	if (preload == NULL || copy == NULL) {
		perror("dtoctl: cannot build LD_PRELOAD");
		exit(EXIT_DTOCTL_FAILURE);
	}
	strcpy(preload, lib);

	for (entry = strtok_r(copy, ": ", &save); entry;
	     entry = strtok_r(NULL, ": ", &save)) {
		if (is_libdto(entry)) {
			real = realpath(entry, NULL);
			if (real == NULL || strcmp(real, lib))
				fprintf(stderr, "dtoctl: dropping '%s' from LD_PRELOAD "
					"in favor of '%s'\n", entry, lib);
			free(real);
			continue;
		}
		strcat(preload, ":");
		strcat(preload, entry);
	}

	if (setenv("LD_PRELOAD", preload, 1)) {
		perror("dtoctl: setenv");
		exit(EXIT_DTOCTL_FAILURE);
	}
	free(copy);
	free(preload);
}

/* Print s, single-quoted if a shell would otherwise split or expand it. */
static void print_quoted(const char *s)
{
	static const char safe[] = "abcdefghijklmnopqrstuvwxyz"
		"ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_@%+=:,./-";

	if (*s && strspn(s, safe) == strlen(s)) {
		fputs(s, stdout);
		return;
	}

	putchar('\'');
	for (; *s; s++) {
		if (*s == '\'')
			fputs("'\\''", stdout);
		else
			putchar(*s);
	}
	putchar('\'');
}

static void print_var(const char *name, int name_len, const char *value,
		      const char *source)
{
	printf("%.*s=", name_len, name);
	print_quoted(value);
	if (source)
		printf("  # %s", source);
	putchar('\n');
}

/* --dry-run: show the library, the environment and the command to be run */
static void print_dry_run(const char *lib, const char *source, char **argv)
{
	char **env;
	int len;

	printf("library: %s (%s)\n", lib, source);
	print_var("LD_PRELOAD", strlen("LD_PRELOAD"), getenv("LD_PRELOAD"), NULL);

	/* variables from options first, in option order, then inherited ones */
	for (int i = 0; i < nr_option_vars; i++)
		print_var(option_vars[i], strlen(option_vars[i]),
			  getenv(option_vars[i]), "option");
	for (env = environ; *env; env++) {
		len = strcspn(*env, "=");
		if (!strncmp(*env, "DTO_", 4) && (*env)[len] == '=' &&
		    !is_option_var(*env, len))
			print_var(*env, len, *env + len + 1, "inherited");
	}

	if (argv == NULL)
		return;
	fputs("command:", stdout);
	for (int i = 0; argv[i]; i++) {
		putchar(' ');
		print_quoted(argv[i]);
	}
	putchar('\n');
}

int main(int argc, char *argv[])
{
	struct opts opts = {0};
	struct argp argp = {
		.options = dtoctl_options,
		.parser = parse_option,
		.args_doc = "<program> [<args>...]",
		.doc = "dtoctl -- run a program with DSA Transparent Offload (DTO) preloaded"
		       "\vOnly the options given are applied: inherited DTO_* variables "
		       "are kept unless an option overrides them or -i is given.  "
		       "Exit status is the program's own, or 125 if dtoctl fails, "
		       "126 if the program cannot be run, 127 if it is not found.",
	};
	const char *source;
	char *lib;
	int err;

	argp_parse(&argp, argc, argv, ARGP_IN_ORDER, NULL, &opts);

	/* clear inherited variables first; this also drops DTO_LIBRARY */
	if (opts.ignore_env)
		clear_dto_environ();
	setup_dto_environ(&opts);

	lib = resolve_library(&opts, &source);
	setup_preload(lib);

	if (opts.dry_run) {
		print_dry_run(lib, source, opts.exename ? argv + opts.idx : NULL);
		return 0;
	}

	/* pass only the target program and its arguments to execvp() */
	execvp(opts.exename, argv + opts.idx);
	err = errno;
	fprintf(stderr, "dtoctl: cannot run '%s': %s\n", opts.exename, strerror(err));

	return err == ENOENT ? EXIT_ENOENT : EXIT_CANNOT_INVOKE;
}