/* A program to put stress on a POSIX system (stress).
 *
 * Copyright 2001-2010 Amos Waterland <apw@rossby.metr.ou.edu>
 * Copyright 2021-2023 Joao Eriberto Mota Filho <eriberto@eriberto.pro.br>
 * Copyright 2023      Vratislav Bendel <vbendel@redhat.com>
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 */

#include <ctype.h>
#include <errno.h>
#include <libgen.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <time.h>
#include <unistd.h>
#include <sys/wait.h>
#include "config.h"

#ifdef HAVE_SYS_PRCTL_H
  #include <sys/prctl.h>
#endif

/* By default, print all messages of severity info and above.  */
static int global_debug = 2;

/* Name of this program */
static char *global_progname = PACKAGE;

/* Implemention of runtime-selectable severity message printing.  */
#define dbg(OUT, STR, ARGS...) if (global_debug >= 3) \
	fprintf (stdout, "%s: dbug: [%lli] ", \
		global_progname, (long long)getpid()), \
		fprintf(OUT, STR, ##ARGS), fflush(OUT)
#define out(OUT, STR, ARGS...) if (global_debug >= 2) \
	fprintf (stdout, "%s: info: [%lli] ", \
		global_progname, (long long)getpid()), \
		fprintf(OUT, STR, ##ARGS), fflush(OUT)
#define wrn(OUT, STR, ARGS...) if (global_debug >= 1) \
	fprintf (stderr, "%s: WARN: [%lli] (%d) ", \
		global_progname, (long long)getpid(), __LINE__), \
		fprintf(OUT, STR, ##ARGS), fflush(OUT)
#define err(OUT, STR, ARGS...) if (global_debug >= 0) \
	fprintf (stderr, "%s: FAIL: [%lli] (%d) ", \
		global_progname, (long long)getpid(), __LINE__), \
		fprintf(OUT, STR, ##ARGS), fflush(OUT)

/* Implementation of check for option argument correctness.  */
#define assert_arg(A) \
          if (++i == argc || ((arg = argv[i])[0] == '-' && \
              !isdigit ((int)arg[1]) )) \
            { \
              err (stderr, "missing argument to option '%s'\n", A); \
              exit (1); \
            }

/* Prototypes for utility functions.  */
int usage (int status);
int version (int status);
long long atoll_s (const char *nptr);
long long atoll_b (const char *nptr);
void worker_init(void);

/* Prototypes for worker functions.  */
int hogcpu (void);
int hogio (void);
int hogvm (long long bytes, long long stride, long long hang, int keep);
int hoghdd (long long bytes);

int
main (int argc, char **argv)
{
    int i, pid, children = 0, retval = 0;
    long starttime, stoptime, runtime, forks;

    /* Variables that indicate which options have been selected.  */
    int do_dryrun = 0;
    long long do_backoff = 3000;
    long long do_timeout = 0;
    long long do_cpu = 0;
    long long do_io = 0;
    long long do_vm = 0;
    long long do_vm_bytes = 256 * 1024 * 1024;
    long long do_vm_stride = 4096;
    long long do_vm_hang = -1;
    int do_vm_keep = 0;
    long long do_hdd = 0;
    long long do_hdd_bytes = 1024 * 1024 * 1024;

    /* Record our start time.  */
    if ((starttime = time (NULL)) == -1)
    {
        err (stderr, "failed to acquire current time: %s\n", strerror (errno));
        exit (1);
    }

    /* SuSv3 does not define any error conditions for this function.  */
    global_progname = basename (argv[0]);

    /* For portability, parse command line options without getopt_long.  */
    for (i = 1; i < argc; i++)
    {
        char *arg = argv[i];

        if (strcmp (arg, "--help") == 0 || strcmp (arg, "-?") == 0)
        {
            usage (0);
        }
        else if (strcmp (arg, "--version") == 0)
        {
            version (0);
        }
        else if (strcmp (arg, "--verbose") == 0 || strcmp (arg, "-v") == 0)
        {
            global_debug = 3;
        }
        else if (strcmp (arg, "--quiet") == 0 || strcmp (arg, "-q") == 0)
        {
            global_debug = 0;
        }
        else if (strcmp (arg, "--dry-run") == 0 || strcmp (arg, "-n") == 0)
        {
            do_dryrun = 1;
        }
        else if (strcmp (arg, "--backoff") == 0)
        {
            assert_arg ("--backoff");
            if (sscanf (arg, "%lli", &do_backoff) != 1)
            {
                err (stderr, "invalid number: %s\n", arg);
                exit (1);
            }
            if (do_backoff < 0)
            {
                err (stderr, "invalid backoff factor: %lli\n", do_backoff);
                exit (1);
            }
            dbg (stdout, "setting backoff coeffient to %llius\n", do_backoff);
        }
        else if (strcmp (arg, "--timeout") == 0 || strcmp (arg, "-t") == 0)
        {
            assert_arg ("--timeout");
            do_timeout = atoll_s (arg);
            if (do_timeout <= 0)
            {
                err (stderr, "invalid timeout value: %llis\n", do_timeout);
                exit (1);
            }
        }
        else if (strcmp (arg, "--cpu") == 0 || strcmp (arg, "-c") == 0)
        {
            assert_arg ("--cpu");
            do_cpu = atoll_b (arg);
            if (do_cpu <= 0)
            {
                err (stderr, "invalid number of cpu hogs: %lli\n", do_cpu);
                exit (1);
            }
        }
        else if (strcmp (arg, "--io") == 0 || strcmp (arg, "-i") == 0)
        {
            assert_arg ("--io");
            do_io = atoll_b (arg);
            if (do_io <= 0)
            {
                err (stderr, "invalid number of io hogs: %lli\n", do_io);
                exit (1);
            }
        }
        else if (strcmp (arg, "--vm") == 0 || strcmp (arg, "-m") == 0)
        {
            assert_arg ("--vm");
            do_vm = atoll_b (arg);
            if (do_vm <= 0)
            {
                err (stderr, "invalid number of vm hogs: %lli\n", do_vm);
                exit (1);
            }
        }
        else if (strcmp (arg, "--vm-bytes") == 0)
        {
            assert_arg ("--vm-bytes");
            do_vm_bytes = atoll_b (arg);
            if (do_vm_bytes <= 0)
            {
                err (stderr, "invalid vm byte value: %lli\n", do_vm_bytes);
                exit (1);
            }
        }
        else if (strcmp (arg, "--vm-stride") == 0)
        {
            assert_arg ("--vm-stride");
            do_vm_stride = atoll_b (arg);
            if (do_vm_stride <= 0)
            {
                err (stderr, "invalid stride value: %lli\n", do_vm_stride);
                exit (1);
            }
        }
        else if (strcmp (arg, "--vm-hang") == 0)
        {
            assert_arg ("--vm-hang");
            do_vm_hang = atoll_b (arg);
            if (do_vm_hang < 0)
            {
                err (stderr, "invalid value: %lli\n", do_vm_hang);
                exit (1);
            }
        }
        else if (strcmp (arg, "--vm-keep") == 0)
        {
            do_vm_keep = 1;
        }
        else if (strcmp (arg, "--hdd") == 0 || strcmp (arg, "-d") == 0)
        {
            assert_arg ("--hdd");
            do_hdd = atoll_b (arg);
            if (do_hdd <= 0)
            {
                err (stderr, "invalid number of hdd hogs: %lli\n", do_hdd);
                exit (1);
            }
        }
        else if (strcmp (arg, "--hdd-bytes") == 0)
        {
            assert_arg ("--hdd-bytes");
            do_hdd_bytes = atoll_b (arg);
            if (do_hdd_bytes <= 0)
            {
                err (stderr, "invalid hdd byte value: %lli\n", do_hdd_bytes);
                exit (1);
            }
        }
        else
        {
            err (stderr, "unrecognized option: %s\n", arg);
            exit (1);
        }
    }

    /* Print startup message if we have work to do, bail otherwise.  */
    if (do_cpu + do_io + do_vm + do_hdd)
    {
        out (stdout, "dispatching hogs: %lli cpu, %lli io, %lli vm, %lli hdd\n",
             do_cpu, do_io, do_vm, do_hdd);
    }
    else
        usage (0);

    /* Round robin dispatch our worker processes.  */
    while ((forks = (do_cpu + do_io + do_vm + do_hdd)))
    {
        long long backoff, timeout = 0;

        /* Calculate the backoff value so we get good fork throughput.  */
        backoff = do_backoff * forks;
        dbg (stdout, "using backoff sleep of %llius\n", backoff);

        /* If we are supposed to respect a timeout, calculate it.  */
        if (do_timeout)
        {
            long long currenttime;

            /* Acquire current time.  */
            if ((currenttime = time (NULL)) == -1)
            {
                perror ("error acquiring current time");
                exit (1);
            }

            /* Calculate timeout based on current time.  */
            timeout = do_timeout - (currenttime - starttime);

            if (timeout > 0)
            {
                dbg (stdout, "setting timeout to %llis\n", timeout);
            }
            else
            {
                wrn (stderr, "used up time before all workers dispatched\n");
                break;
            }
        }

        if (do_cpu)
        {
            switch (pid = fork ())
            {
            case 0:            /* child */
                worker_init();
                alarm (timeout);
                usleep (backoff);
                if (do_dryrun)
                    exit (0);
                exit (hogcpu ());
            case -1:           /* error */
                err (stderr, "fork failed: %s\n", strerror (errno));
                break;
            default:           /* parent */
                dbg (stdout, "--> hogcpu worker %lli [%i] forked\n",
                     do_cpu, pid);
                ++children;
            }
            --do_cpu;
        }

        if (do_io)
        {
            switch (pid = fork ())
            {
            case 0:            /* child */
                worker_init();
                alarm (timeout);
                usleep (backoff);
                if (do_dryrun)
                    exit (0);
                exit (hogio ());
            case -1:           /* error */
                err (stderr, "fork failed: %s\n", strerror (errno));
                break;
            default:           /* parent */
                dbg (stdout, "--> hogio worker %lli [%i] forked\n", do_io, pid);
                ++children;
            }
            --do_io;
        }

        if (do_vm)
        {
            switch (pid = fork ())
            {
            case 0:            /* child */
                worker_init();
                alarm (timeout);
                usleep (backoff);
                if (do_dryrun)
                    exit (0);
                exit (hogvm (do_vm_bytes, do_vm_stride, do_vm_hang, do_vm_keep));
            case -1:           /* error */
                err (stderr, "fork failed: %s\n", strerror (errno));
                break;
            default:           /* parent */
                dbg (stdout, "--> hogvm worker %lli [%i] forked\n", do_vm, pid);
                ++children;
            }
            --do_vm;
        }

        if (do_hdd)
        {
            switch (pid = fork ())
            {
            case 0:            /* child */
                worker_init();
                alarm (timeout);
                usleep (backoff);
                if (do_dryrun)
                    exit (0);
                exit (hoghdd (do_hdd_bytes));
            case -1:           /* error */
                err (stderr, "fork failed: %s\n", strerror (errno));
                break;
            default:           /* parent */
                dbg (stdout, "--> hoghdd worker %lli [%i] forked\n",
                     do_hdd, pid);
                ++children;
            }
            --do_hdd;
        }
    }

    /* Wait for our children to exit.  */
    while (children)
    {
        int status, ret;

        if ((pid = wait (&status)) > 0)
        {
            --children;

            if (WIFEXITED (status))
            {
                if ((ret = WEXITSTATUS (status)) == 0)
                {
                    dbg (stdout, "<-- worker %i returned normally\n", pid);
                }
                else
                {
                    err (stderr, "<-- worker %i returned error %i\n", pid, ret);
                    ++retval;
                    wrn (stderr, "now reaping child worker processes\n");
                    if (signal (SIGUSR1, SIG_IGN) == SIG_ERR)
                        err (stderr, "handler error: %s\n", strerror (errno));
                    if (kill (-1 * getpid (), SIGUSR1) == -1)
                        err (stderr, "kill error: %s\n", strerror (errno));
                }
            }
            else if (WIFSIGNALED (status))
            {
                if ((ret = WTERMSIG (status)) == SIGALRM)
                {
                    dbg (stdout, "<-- worker %i signalled normally\n", pid);
                }
                else if ((ret = WTERMSIG (status)) == SIGUSR1)
                {
                    dbg (stdout, "<-- worker %i reaped\n", pid);
                }
                else
                {
                    err (stderr, "<-- worker %i got signal %i\n", pid, ret);
                    ++retval;
                    wrn (stderr, "now reaping child worker processes\n");
                    if (signal (SIGUSR1, SIG_IGN) == SIG_ERR)
                        err (stderr, "handler error: %s\n", strerror (errno));
                    if (kill (-1 * getpid (), SIGUSR1) == -1)
                        err (stderr, "kill error: %s\n", strerror (errno));
                }
            }
            else
            {
                err (stderr, "<-- worker %i exited abnormally\n", pid);
                ++retval;
            }
        }
        else
        {
            err (stderr, "error waiting for worker: %s\n", strerror (errno));
            ++retval;
            break;
        }
    }

    /* Record our stop time.  */
    if ((stoptime = time (NULL)) == -1)
    {
        err (stderr, "failed to acquire current time\n");
        exit (1);
    }

    /* Calculate our runtime.  */
    runtime = stoptime - starttime;

    /* Print final status message.  */
    if (retval)
    {
        err (stderr, "failed run completed in %lis\n", runtime);
    }
    else
    {
        out (stdout, "successful run completed in %lis\n", runtime);
    }

    exit (retval);
}

void worker_init(void)
{
#ifdef HAVE_SYS_PRCTL_H
    if (prctl(PR_SET_PDEATHSIG, SIGTERM) == -1)
        exit(0);
#endif
}

int
hogcpu (void)
{
    while (1)
        sqrt (rand ());

    return 0;
}

int
hogio ()
{
    while (1)
        sync ();

    return 0;
}

int
hogvm (long long bytes, long long stride, long long hang, int keep)
{
    long long i;
    char *ptr = 0;
    char c;
    int do_malloc = 1;

    while (1)
    {
        if (do_malloc)
        {
            dbg (stdout, "allocating %lli bytes ...\n", bytes);
            if (!(ptr = (char *) malloc (bytes * sizeof (char))))
            {
                err (stderr, "hogvm malloc failed: %s\n", strerror (errno));
                return 1;
            }
            if (keep)
                do_malloc = 0;
        }

        dbg (stdout, "touching bytes in strides of %lli bytes ...\n", stride);
        for (i = 0; i < bytes; i += stride)
            ptr[i] = 'Z';           /* Ensure that COW happens.  */

        if (hang == 0)
        {
            dbg (stdout, "sleeping forever with allocated memory\n");
            while (1)
                sleep (1024);
        }
        else if (hang > 0)
        {
            dbg (stdout, "sleeping for %llis with allocated memory\n", hang);
            sleep (hang);
        }

        for (i = 0; i < bytes; i += stride)
        {
            c = ptr[i];
            if (c != 'Z')
            {
                err (stderr, "memory corruption at: %p\n", ptr + i);
                return 1;
            }
        }

        if (do_malloc)
        {
            free (ptr);
            dbg (stdout, "freed %lli bytes\n", bytes);
        }
    }

    return 0;
}

int
hoghdd (long long bytes)
{
    long long i, j;
    int fd;
    int chunk = (1024 * 1024) - 1;        /* Minimize slow writing.  */
    char buff[chunk];

    /* Initialize buffer with some random ASCII data.  */
    dbg (stdout, "seeding %d byte buffer with random data\n", chunk);
    for (i = 0; i < chunk - 1; i++)
    {
        j = rand ();
        j = (j < 0) ? -j : j;
        j %= 95;
        j += 32;
        buff[i] = j;
    }
    buff[i] = '\n';

    while (1)
    {
        char name[] = "./stress.XXXXXX";

        if ((fd = mkstemp (name)) == -1)
        {
            err (stderr, "mkstemp failed: %s\n", strerror (errno));
            return 1;
        }

        dbg (stdout, "opened %s for writing %lli bytes\n", name, bytes);

        dbg (stdout, "unlinking %s\n", name);
        if (unlink (name) == -1)
        {
            err (stderr, "unlink of %s failed: %s\n", name, strerror (errno));
            return 1;
        }

        dbg (stdout, "fast writing to %s\n", name);
        for (j = 0; bytes == 0 || j + chunk < bytes; j += chunk)
        {
            if (write (fd, buff, chunk) == -1)
            {
                err (stderr, "write failed: %s\n", strerror (errno));
                return 1;
            }
        }

        dbg (stdout, "slow writing to %s\n", name);
        for (; bytes == 0 || j < bytes - 1; j++)
        {
            if (write (fd, &buff[j % chunk], 1) == -1)
            {
                err (stderr, "write failed: %s\n", strerror (errno));
                return 1;
            }
        }
        if (write (fd, "\n", 1) == -1)
        {
            err (stderr, "write failed: %s\n", strerror (errno));
            return 1;
        }
        ++j;

        dbg (stdout, "closing %s after %lli bytes\n", name, j);
        close (fd);
    }

    return 0;
}

/* Convert a string representation of a number with an optional size suffix
 * to a long long.
 */
long long
atoll_b (const char *nptr)
{
    int pos;
    char suffix;
    long long factor = 0;
    long long value;

    if ((pos = strlen (nptr) - 1) < 0)
    {
        err (stderr, "invalid string\n");
        exit (1);
    }

    switch (suffix = nptr[pos])
    {
    case 'b':
    case 'B':
        factor = 0;
        break;
    case 'k':
    case 'K':
        factor = 10;
        break;
    case 'm':
    case 'M':
        factor = 20;
        break;
    case 'g':
    case 'G':
        factor = 30;
        break;
    default:
        if (suffix < '0' || suffix > '9')
        {
            err (stderr, "unrecognized suffix: %c\n", suffix);
            exit (1);
        }
    }

    if (sscanf (nptr, "%lli", &value) != 1)
    {
        err (stderr, "invalid number: %s\n", nptr);
        exit (1);
    }

    value = value << factor;

    return value;
}

/* Convert a string representation of a number with an optional time suffix
 * to a long long.
 */
long long
atoll_s (const char *nptr)
{
    int pos;
    char suffix;
    long long factor = 1;
    long long value;

    if ((pos = strlen (nptr) - 1) < 0)
    {
        err (stderr, "invalid string\n");
        exit (1);
    }

    switch (suffix = nptr[pos])
    {
    case 's':
    case 'S':
        factor = 1;
        break;
    case 'm':
    case 'M':
        factor = 60;
        break;
    case 'h':
    case 'H':
        factor = 60 * 60;
        break;
    case 'd':
    case 'D':
        factor = 60 * 60 * 24;
        break;
    case 'y':
    case 'Y':
        factor = 60 * 60 * 24 * 365;
        break;
    default:
        if (suffix < '0' || suffix > '9')
        {
            err (stderr, "unrecognized suffix: %c\n", suffix);
            exit (1);
        }
    }

    if (sscanf (nptr, "%lli", &value) != 1)
    {
        err (stderr, "invalid number: %s\n", nptr);
        exit (1);
    }

    value = value * factor;

    return value;
}

int
version (int status)
{
    char *mesg = "%s %s\n";

    fprintf (stdout, mesg, global_progname, VERSION);

    if (status <= 0)
        exit (-1 * status);

    return 0;
}

int
usage (int status)
{
    char *mesg =
        "`%s' imposes certain types of compute stress on your system\n\n"
        "Usage: %s [OPTION [ARG]] ...\n"
        " -?, --help         show this help statement\n"
        "     --version      show version statement\n"
        " -v, --verbose      be verbose\n"
        " -q, --quiet        be quiet\n"
        " -n, --dry-run      show what would have been done\n"
        " -t, --timeout N    timeout after N seconds\n"
        "     --backoff N    wait factor of N microseconds before work starts\n"
        " -c, --cpu N        spawn N workers spinning on sqrt()\n"
        " -i, --io N         spawn N workers spinning on sync()\n"
        " -m, --vm N         spawn N workers spinning on malloc()/free()\n"
        "     --vm-bytes B   malloc B bytes per vm worker (default is 256MB)\n"
        "     --vm-stride B  touch a byte every B bytes (default is 4096)\n"
        "     --vm-hang N    sleep N secs before free (default none, 0 is inf)\n"
        "     --vm-keep      redirty memory instead of freeing and reallocating\n"
        " -d, --hdd N        spawn N workers spinning on write()/unlink()\n"
        "     --hdd-bytes B  write B bytes per hdd worker (default is 1GB)\n\n"
        "Example: %s --cpu 8 --io 4 --vm 2 --vm-bytes 128M --timeout 10s\n\n"
        "Note: Numbers may be suffixed with s,m,h,d,y (time) or B,K,M,G (size).\n";

    fprintf (stdout, mesg, global_progname, global_progname, global_progname);

    if (status <= 0)
        exit (-1 * status);

    return 0;
}
