/*
 * Copyright (c) 2004, 2022 Roger Seguin <rgrsgn@gmail.com, https://frogzie.duckdns.org>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

static char     rcsid[] =
    "$Id: diskiobench.C,v 2.1 2023/06/07 10:36:15 RogerSeguin Exp $"
    "\nCopyright (c) 2004, 2022 Roger Seguin <rgrsgn@gmail.com"
    " https://frogzie.duckdns.org>";

/*
Description:
	Indicate how much time it takes for reading/writing disk data
	(disk I/O benchmark)
	Usage:	diskiobench <data size in MiB>
*/

#ifndef _GNU_SOURCE
#define _GNU_SOURCE	1
#endif

#include <unistd.h>
#include <stdlib.h>
#include <fcntl.h>
#include <memory.h>
#include <strings.h>
#include <string.h>
#include <time.h>
#include <limits.h>
#include <stdio.h>
#include <signal.h>
#include <sys/stat.h>
#include <sys/time.h>


#define TMPFILE	"diskiobench.dat"


namespace {
    char *m_pcTmpFile = 0;
};


static int Usage (const char *const This, FILE * p_pFOut)
{
    FILE           *pFOut = p_pFOut;
    fprintf (stdout, "%s\n", rcsid + 5);
    fprintf (stdout, "Disk benchmarking tool\n");
    fprintf (pFOut, "Usage\n");
    fprintf (pFOut, "\t%s [-b blockSize_kiB] [-f tmpfile] [sizeMiB]\n",
             This);
    fprintf (pFOut, "\t%s -h\n", This);
    fprintf (pFOut, "Options\n");
    fprintf (pFOut,
             "\t\t-b blockSize:\tread/write by blocks of indicated size\n"
             "\t\t\t\tin kibibytes (default is 64 kiB)\n");
    fprintf (pFOut, "\t\t-f file:\t"
             "temporary file used for disk benchmarking\n"
             "\t\t\t\t(default=\"%s\")\n", TMPFILE);
    fprintf (pFOut, "\t\t-h:\t\toutput this help and exit\n");
    fprintf (pFOut, "Description\n");
    fprintf (pFOut,
             "\tTime write/read the benchmarking file tmpfile (defaul size: 1024 MiB)\n");
    fprintf (pFOut,
             "\t(If the indicated source is a directory, \"" TMPFILE "\" will be used)\n");
    fprintf (pFOut,
             "\ttmpfile is removed after measurement\n");
    return (0);
}


static int64_t CurrentTime_ms ()
{
    struct timespec  oTime;
    int             status;

    status = clock_gettime (CLOCK_REALTIME, &oTime);
    if (status == -1) {
        perror ("clock_gettime()");
        return (-1);
    }
    return (oTime.tv_sec * 1000 + oTime.tv_nsec / 1000000);
}


static int Write (const size_t p_DataSize_MiB, const char *p_FileName,
                  const size_t p_BlockSize_B)
{
    const int64_t   DataSize_B = (int64_t) p_DataSize_MiB * 1024 * 1024;
    const int       nloops = DataSize_B / p_BlockSize_B;
    char            acBuffer[p_BlockSize_B];
    int             fd;
    int64_t	    nbytes = 0;
    int64_t         t1, t2;	// ms

    memset (acBuffer, 0, sizeof (acBuffer));
    fd = open("/dev/urandom", O_RDONLY);
    if (fd == -1)
        perror ("/dev/urandom");
    else {
        read(fd, acBuffer, sizeof (acBuffer));
        close(fd);
    }

    t1 = CurrentTime_ms ();
    fd = open (p_FileName, O_CREAT | O_RDWR, S_IRUSR + S_IWUSR);
    if (fd == -1) {
        perror (p_FileName);
        return (-1);
    }
    for (int i = 0; i < nloops; i++) {
        ssize_t n =  write (fd, acBuffer, sizeof (acBuffer));
        if (n == -1) {
            perror ("write()");
            close (fd);
            return (-1);
        }
        else if (!n)
            break;
        nbytes += n;
    }
    fsync (fd);
    close (fd);
    t2 = CurrentTime_ms ();
    if (nbytes != DataSize_B) {
        fprintf (stderr, "Unknown writing error\n");
        return (-1);
    }
    return (static_cast < int >(t2 - t1));
}


int Read (const size_t p_DataSize_MiB, const char *p_FileName,
          const size_t p_BlockSize_B)
{
    const int64_t   DataSize_B = (int64_t) p_DataSize_MiB * 1024 * 1024;
    const int       nloops = DataSize_B / p_BlockSize_B;
    char            acBuffer[p_BlockSize_B];
    int             fd;
    int64_t	    nbytes = 0;
    int64_t         t1, t2;	// ms

    t1 = CurrentTime_ms ();
    fd = open (p_FileName, O_RDONLY);
    if (fd == -1) {
        perror (p_FileName);
        return (-1);
    }
    for (int i = 0; i < nloops; i++) {
        ssize_t n = read (fd, acBuffer, sizeof (acBuffer));
        if (n == -1) {
            perror ("read()");
            close (fd);
            return (-1);
        }
        else if (!n)
            break;
        nbytes += n;
    }
    close (fd);
    t2 = CurrentTime_ms ();
    if (nbytes != DataSize_B) {
        fprintf (stderr, "Unknown reading error\n");
        return (-1);
    }
    return (static_cast < int >(t2 - t1));
}


static void printResult(const char *p_pcAction, size_t p_DataSize_MiB,
                        int p_duration_ms, size_t p_BlockSize_kiB)
{
    size_t speedMiB_perSec
        = (p_duration_ms ? (1000 * p_DataSize_MiB / p_duration_ms) : 0);
    size_t speedMB_perSec = speedMiB_perSec * 1024 * 1024 / 1000000;
    double speedGbit_perSec = 8.0 * speedMB_perSec / 1000.0;
    printf (
        "%s\t%lu MiB by %lu-kiB blocks in %u ms (%lu MiB/s = %lu MB/s = %.1f Gbit/s)\n",
        p_pcAction, p_DataSize_MiB, p_BlockSize_kiB, p_duration_ms,
        speedMiB_perSec, speedMB_perSec, speedGbit_perSec);
}


static bool isDirectory(const char* path)
{
    struct stat attr;
    int status = stat(path, &attr);
    return ((status == 0) && S_ISDIR(attr.st_mode));
}


static void setFilename (char *filename,
                         const char *path, const char *defaultBasename)
{
    strcpy (filename, path);
    char *pc = filename + strlen(filename) - 1;
    if (*pc == '/')
        *pc = 0;	/* Remove trailing "/" */
    if (isDirectory (filename)) {
        strcat (filename, "/");
        strcat (filename, defaultBasename);
    }
}


static int resetCache()
{
    int uid = getuid();
    if (uid != 0)
        return (1);
    const char cache[] = "/proc/sys/vm/drop_caches";
    int fd = open (cache, O_WRONLY);
    if (fd == -1) {
        perror (cache);
        return (-1);
    }
    char drop[] = "1\n";
    write (fd, drop, sizeof (drop));
    close (fd);
    return (0);
}


static void interrupt (int p_iSignal, siginfo_t *, void *)
/* Signal handler */
{
    if (m_pcTmpFile)
        unlink (m_pcTmpFile);
    exit (p_iSignal);
}


int main (int argc, char **argv)
{
    size_t          DataSize_MiB = 1024;	/* How much to read/write? */
    static char     FileName[PATH_MAX] = TMPFILE;	/* where? */
    size_t          BlockSize_kiB = 16;	/* by blocks of what size? */
    char            c;
    int             fError = 0, fUsage = 0;
    int             d;	// ms
    int64_t         t1, t2;	// ms
    struct sigaction oSigaction;
    int             status;

    /* Parse command line options */
    putenv ((char *) "POSIXLY_CORRECT=1");
    for (;
            !fUsage && !fError
            && (c = getopt (argc, argv, "b:f:h")) != EOF;)
        switch (c) {
            case 'b':	/* Block size */
                BlockSize_kiB = atoi (optarg);
                break;
            case 'f':	/* Name and location of temporary file to read/write */
                setFilename (FileName, optarg, TMPFILE);
                break;
            case 'h':	/* Help */
                fUsage = 1;
                break;
            case '?':
            default:
                fError = 1;
        }
    if (!fUsage) {
        fError |= (argc > (optind + 1));
        if (!fError) {
            if (argc == (optind + 1))
                DataSize_MiB = atoi (argv[optind]);
            fError = (DataSize_MiB <= 0);
        }
    }
    fUsage |= fError;

    /* Display usage message whether requested or error detected */
    if (fUsage) {
        Usage (basename (argv[0]), fError ? stderr : stdout);
        return (fError ? ~0 : 0);
    }

    int uid = getuid();
    if (uid != 0)
        printf (
            " ** Warning: this program should be run as \"root\" for better accuracy\n"
        );

    m_pcTmpFile = FileName;
    memset (&oSigaction, 0, sizeof (oSigaction));
    oSigaction.sa_sigaction = interrupt;
    sigaction (SIGHUP, &oSigaction, 0);
    sigaction (SIGINT, &oSigaction, 0);
    sigaction (SIGTERM, &oSigaction, 0);

    printf ("Using %s\n", FileName);
    unlink (FileName);
    sync();

    size_t BlockSize_B = BlockSize_kiB * 1024;

    d = Write (DataSize_MiB, FileName, BlockSize_B);
    if (d == -1)
        return (~0);
    printResult("Wrote", DataSize_MiB, d, BlockSize_kiB);

    resetCache();
    d = Read (DataSize_MiB, FileName, BlockSize_B);
    if (d != -1)
        printResult("Read", DataSize_MiB, d, BlockSize_kiB);

    d = Write (DataSize_MiB, FileName, BlockSize_B);
    if (d != -1)
        printResult("Rewrote", DataSize_MiB, d, BlockSize_kiB);

    t1 = CurrentTime_ms ();
    status = unlink (FileName);
    if (status == -1)
        perror ("unlink()");
    t2 = CurrentTime_ms ();
    d = (int) (t2 - t1);
    printf ("Deleted\t%lu MiB in %d ms\n", DataSize_MiB, d);

    return (0);
}


/*
 * $Log: diskiobench.C,v $
 * Revision 2.1  2023/06/07 10:36:15  RogerSeguin
 * Clear memory cache when run as root
 *
 * Revision 2.0  2022/10/24 08:47:07  RogerSeguin
 * Parameterizable block size
 *
 * Revision 1.1  2004/07/26 12:45:16  RogerSeguin
 * Initial revision
 *
*/