/*
 *  Testing flash card global wear leveling under stress
 *
 *  Copyright (C) 2011  Clifford Wolf <clifford@clifford.at>
 *
 *  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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 *  Compile and run:
 *  gcc -std=gnu99 -Wall -Werror -Os -pthread -o flashcrash flashcrash.c -lm
 *  ./flashcrash init /dev/<diskdev>
 *  ./flashcrash crash /dev/<diskdev>   (Ctrl-C after some hours)
 *  ./flashcrash check /dev/<diskdev>
 *
 *  Or:
 *  ./flashcrash init-and-crash /dev/<diskdev>   (Ctrl-C after some hours)
 *  ./flashcrash check /dev/<diskdev>
 *
 *  Or:
 *  ./flashcrash auto /dev/<diskdev>
 *
 *  Some background information:
 *  Some flash drives with global waer leveling loose data when transfering
 *  blocks under heavy stress. This program can be used to provoke this by
 *  writing the same logical hd blocks very often and so triggering many
 *  transfers of nevwer-written blocks. Do not use this tool on flash cards
 *  with dynamic waer leveling. They won't last very long..
 */

#define _GNU_SOURCE
#define _LARGEFILE64_SOURCE

#include <stdio.h>
#include <errno.h>
#include <string.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdint.h>
#include <stdbool.h>
#include <math.h>
#include <time.h>
#include <signal.h>
#include <pthread.h>

/**** BEGIN: http://svn.clifford.at/tools/trunk/examples/check.h ****/

// This is to not confuse the VIM syntax highlighting
#define CHECK_VAL_OPEN (
#define CHECK_VAL_CLOSE )

#define CHECK(result, check)                                         \
   CHECK_VAL_OPEN{                                                   \
     typeof(result) _R = (result);                                   \
     if (!(_R check)) {                                              \
       fprintf(stderr, "Error from '%s' (%ld %s) in %s:%d.\n",       \
               #result, (long int)_R, #check, __FILE__, __LINE__);   \
       fprintf(stderr, "ERRNO(%d): %s\n", errno, strerror(errno));   \
       abort();                                                      \
     }                                                               \
     _R;                                                             \
   }CHECK_VAL_CLOSE

/**** END: http://svn.clifford.at/tools/trunk/examples/check.h ****/


#define NUM_BUFFERS 4
#define BLOCK_SIZE (16*1024*1024)

int opt_scratchpad_pos = 0;
int opt_reuse_buffers = 0;
int opt_autocycles = 1;

int fd;
off64_t dev_size;

bool got_ctrl_c;
pthread_mutex_t total_written_mutex = PTHREAD_MUTEX_INITIALIZER;
off64_t total_written;

pthread_mutex_t buf_mutex = PTHREAD_MUTEX_INITIALIZER;
pthread_cond_t buf_cond_prep = PTHREAD_COND_INITIALIZER;
pthread_cond_t buf_cond_write = PTHREAD_COND_INITIALIZER;
bool buffer_busy[NUM_BUFFERS];
int buffer_usecount[NUM_BUFFERS];
char buffers[NUM_BUFFERS][BLOCK_SIZE] __attribute__ ((aligned (512)));
pthread_t th_prep, th_write;

void do_flush()
{
	printf("Flushing all caches..\n");
	sync();
	FILE *fp = CHECK(fopen("/proc/sys/vm/drop_caches", "wt"), != NULL);
	fprintf(fp, "3\n");
	fclose(fp);
}

void do_open(const char *dev, int mode)
{
	do_flush();
	fd = CHECK(open(dev, mode), >= 0);
	dev_size = lseek64(fd, 0, SEEK_END);
	printf("Opened %s (size = %.3f GB).\n", dev, dev_size / pow(1024, 3));
}

void fillbuf(char *buffer, uint32_t num)
{
	if (num == 0) {
		memset(buffer, 0, BLOCK_SIZE);
		return;
	}

	uint32_t d = num;
	for (int p=0; p<BLOCK_SIZE; p+=4) {
		// LCG as recommended in Numerical Recipes
		// see also http://en.wikipedia.org/wiki/Linear_congruential_generator
		d = 1664525 * d + 1013904223;
		buffer[p+0] = d;
		buffer[p+1] = d >> 8;
		buffer[p+2] = d >> 16;
		buffer[p+3] = d >> 24;
	}
}

void do_init()
{
	time_t start = time(NULL);
	time_t laststat = 0;
	lseek64(fd, 0, SEEK_SET);
	printf("Performing init with a blocksize of %.2f MB..\n", BLOCK_SIZE / pow(1024, 2));
	for (int n=0; n < dev_size / BLOCK_SIZE; n++)
	{
		if (laststat != time(NULL) || n == (dev_size / BLOCK_SIZE) - 1) {
			laststat = time(NULL);
			double gb_written = n * (BLOCK_SIZE / pow(1024, 3));
			double gb_total = dev_size / pow(1024, 3);
			double percent_done = 100 * gb_written / gb_total;
			double mb_per_s = 1024 * gb_written / (0.1 + time(NULL) - start);
			int eta_tot_sec = 1024 * (gb_total - gb_written) / (mb_per_s + 0.001);
			int eta_hrs = eta_tot_sec / (60*60) < 99 ? eta_tot_sec / (60*60) : 99;
			int eta_min = (eta_tot_sec / 60) % 60, eta_sec = eta_tot_sec % 60;
			printf("init: %6.3f GB / %6.3f GB      %6.2f%%, %6.2f MB/s,    ETA: %02d:%02d:%02d    \r",
					gb_written, gb_total, percent_done, mb_per_s, eta_hrs, eta_min, eta_sec);
			fflush(stdout);
		}
		fillbuf(buffers[0], n);
		CHECK(write(fd, buffers[0], BLOCK_SIZE), == BLOCK_SIZE);
	}
	printf("\n");
}

void ctrl_c_hdlr(int unused)
{
	got_ctrl_c = true;
}

void *do_crash_prep_worker(void *arg)
{
	int count = 0;
	struct timeval now;
	struct timespec timeout;
	pthread_mutex_lock(&buf_mutex);
	while (!got_ctrl_c)
	{
		for (int bufnum=0; bufnum < NUM_BUFFERS; bufnum++)
		{
			if (buffer_busy[bufnum])
				continue;
			if (buffer_usecount[bufnum] <= opt_reuse_buffers/2)
				continue;

			buffer_busy[bufnum] = true;
			pthread_mutex_unlock(&buf_mutex);

			fillbuf(buffers[bufnum], (count++) | 0x80000000);
			buffer_usecount[bufnum] = 0;

			pthread_mutex_lock(&buf_mutex);
			buffer_busy[bufnum] = false;
			pthread_cond_signal(&buf_cond_write);
		}
		gettimeofday(&now, NULL);
		timeout.tv_sec = now.tv_sec + 5;
		timeout.tv_nsec = now.tv_usec * 1000;
		pthread_cond_timedwait(&buf_cond_prep, &buf_mutex, &timeout);
	}
	pthread_mutex_unlock(&buf_mutex);
	return NULL;
}

void *do_crash_write_worker(void *arg)
{
	int bufnum = 0;
	int lastbufnum = -1;
	struct timeval now;
	struct timespec timeout;
	pthread_mutex_lock(&buf_mutex);
	while (!got_ctrl_c)
	{
		for (int i=0; i < NUM_BUFFERS; i++) {
			bufnum = bufnum < NUM_BUFFERS-1 ? bufnum+1 : 0;
			if (bufnum == lastbufnum)
				continue;
			if (buffer_usecount[bufnum] > opt_reuse_buffers)
				continue;
			goto found_next_buffer;
		}

		printf("\nWaiting for buffers to be prepared. => Writespeed is CPU bound!\n");
		gettimeofday(&now, NULL);
		timeout.tv_sec = now.tv_sec + 5;
		timeout.tv_nsec = now.tv_usec * 1000;
		pthread_cond_timedwait(&buf_cond_write, &buf_mutex, &timeout);
		continue;

found_next_buffer:
		buffer_busy[bufnum] = true;
		pthread_mutex_unlock(&buf_mutex);

		lseek64(fd, opt_scratchpad_pos * (off64_t)BLOCK_SIZE, SEEK_SET);
		CHECK(write(fd, buffers[bufnum], BLOCK_SIZE), == BLOCK_SIZE);

		pthread_mutex_lock(&total_written_mutex);
		total_written += BLOCK_SIZE;
		pthread_mutex_unlock(&total_written_mutex);

		pthread_mutex_lock(&buf_mutex);
		buffer_busy[bufnum] = false;
		buffer_usecount[bufnum]++;
		lastbufnum = bufnum;

		pthread_cond_signal(&buf_cond_prep);
	}
	pthread_mutex_unlock(&buf_mutex);
	return NULL;
}

enum crash_status_e {
	crash_status_gotctrlc = 1,
	crash_status_reachedlimit = 2,
};

enum crash_status_e do_crash(off64_t limit)
{
	enum crash_status_e rc = crash_status_gotctrlc;
	bool reached_limit = false;

	got_ctrl_c = false;
	total_written = 0;

	signal(SIGINT, ctrl_c_hdlr);
	printf("Performing crash procedure with a blocksize of %.2f MB..\n", BLOCK_SIZE / pow(1024, 2));

	for (int i=0; i<NUM_BUFFERS; i++) {
		buffer_busy[i] = false;
		buffer_usecount[i] = opt_reuse_buffers + 1;
	}

	pthread_create(&th_prep, NULL, do_crash_prep_worker, NULL);

	pthread_mutex_lock(&buf_mutex);
	for (int i=0; i < NUM_BUFFERS; i++) {
		while (buffer_usecount[i] > opt_reuse_buffers)
			pthread_cond_wait(&buf_cond_write, &buf_mutex);
	}
	pthread_mutex_unlock(&buf_mutex);

	time_t start = time(NULL);
	pthread_create(&th_write, NULL, do_crash_write_worker, NULL);

	while (!got_ctrl_c && !reached_limit) {
		sleep(1);
		pthread_mutex_lock(&total_written_mutex);
		double gb_written = total_written / pow(1024, 3);
		if (limit != 0 && total_written >= limit)
			reached_limit = true;
		pthread_mutex_unlock(&total_written_mutex);
		double mb_per_s = 1024 * gb_written / (0.1 + time(NULL) - start);
		if (limit != 0) {
			double gb_total = limit / pow(1024, 3);
			double percent_done = 100 * gb_written / gb_total;
			int eta_tot_sec = 1024 * (gb_total - gb_written) / (mb_per_s + 0.001);
			int eta_hrs = eta_tot_sec / (60*60) < 99 ? eta_tot_sec / (60*60) : 99;
			int eta_min = (eta_tot_sec / 60) % 60, eta_sec = eta_tot_sec % 60;
			printf("crash: %6.3f GB / %6.3f GB  (iter count = %.0f)  %6.2f%%, %6.2f MB/s,  ETA: %02d:%02d:%02d  <",
					gb_written, gb_total, gb_written * pow(1024, 3) / BLOCK_SIZE,
					percent_done, mb_per_s, eta_hrs, eta_min, eta_sec);
		} else {
			printf("crash: %6.3f GB  (iter count = %.0f)  %6.2f MB/s  <",
					gb_written, gb_written * pow(1024, 3) / BLOCK_SIZE, mb_per_s);
		}
		for (int i=0; i < NUM_BUFFERS; i++)
			printf("%3d", buffer_usecount[i]);
		printf(" >   \r");
		fflush(stdout);
	}

	if (reached_limit) {
		printf("\nReached size limit for crash iteration. Waiting for worker threads...\n");
		rc = crash_status_reachedlimit;
		got_ctrl_c = true;
	} else
		printf("\nGot Ctrl-C. Waiting for worker threads...\n");

	pthread_join(th_write, NULL);
	pthread_join(th_prep, NULL);

	fillbuf(buffers[0], opt_scratchpad_pos);
	lseek64(fd, opt_scratchpad_pos * (off64_t)BLOCK_SIZE, SEEK_SET);
	CHECK(write(fd, buffers[0], BLOCK_SIZE), == BLOCK_SIZE);

	signal(SIGINT, SIG_DFL);
	return rc;
}

bool do_check()
{
	bool found_errors = false;
	time_t start = time(NULL);
	time_t laststat = 0;
	lseek64(fd, 0, SEEK_SET);
	printf("Performing check with a blocksize of %.2f MB..\n", BLOCK_SIZE / pow(1024, 2));
	for (int n=0; n < dev_size / BLOCK_SIZE; n++)
	{
		fillbuf(buffers[0], n);
		CHECK(read(fd, buffers[1], BLOCK_SIZE), == BLOCK_SIZE);
		if (memcmp(buffers[0], buffers[1], BLOCK_SIZE)) {
			int biterrors = 0;
			for (int i=0; i < BLOCK_SIZE*8; i++) {
				int bit1 = (buffers[0][i/8] >> (i%8)) & 1;
				int bit2 = (buffers[1][i/8] >> (i%8)) & 1;
				if (bit1 != bit2)
					biterrors++;
			}
			printf("Found error in block %d (off = %Ld): %d of %d bits (%.2f%%) don't match.\n",
					n, n * (long long)BLOCK_SIZE, biterrors, BLOCK_SIZE*8,
					biterrors * (100.0 / (BLOCK_SIZE*8)));
			found_errors = true;
			goto print_check_progress;
		}
		if (laststat != time(NULL) || n == (dev_size / BLOCK_SIZE) - 1) {
print_check_progress:;
			laststat = time(NULL);
			double gb_written = n * (BLOCK_SIZE / pow(1024, 3));
			double gb_total = dev_size / pow(1024, 3);
			double percent_done = 100 * gb_written / gb_total;
			double mb_per_s = 1024 * gb_written / (0.1 + time(NULL) - start);
			int eta_tot_sec = 1024 * (gb_total - gb_written) / (mb_per_s + 0.001);
			int eta_hrs = eta_tot_sec / (60*60) < 99 ? eta_tot_sec / (60*60) : 99;
			int eta_min = (eta_tot_sec / 60) % 60, eta_sec = eta_tot_sec % 60;
			printf("check: %6.3f GB / %6.3f GB      %6.2f%%, %6.2f MB/s,    ETA: %02d:%02d:%02d    \r",
					gb_written, gb_total, percent_done, mb_per_s, eta_hrs, eta_min, eta_sec);
			fflush(stdout);
		}
	}
	printf("\n");
	return found_errors;
}

void do_close(char *what)
{
	do_flush();
	close(fd);
	printf("%s done. Device closed.\n", what);
}

int main(int argc, char **argv)
{
	char *progname = argc >= 1 ? argv[0] : "flashcrash";
	int opt;

	while ((opt = getopt(argc, argv, "p:r:c:")) != -1)
	{
		switch (opt)
		{
		case 'p':
			opt_scratchpad_pos = atoi(optarg);
			printf("Scratchpad position at block %d (at %.2f GB)\n",
					opt_scratchpad_pos, opt_scratchpad_pos * (BLOCK_SIZE / pow(1024,3)));
			break;
		case 'r':
			opt_reuse_buffers = atoi(optarg);
			printf("Reusing each prepped crash buffer %d times.\n", opt_reuse_buffers);
			break;
		case 'c':
			opt_autocycles = atoi(optarg);
			printf("Number of crash cycles per check cycle: %d\n", opt_autocycles);
			break;
		default:
			goto print_usage;
		}
	}

	argc -= (optind-1);
	argv += (optind-1);

	if (argc == 3 && !strcmp(argv[1], "init")) {
		do_open(argv[2], O_WRONLY|O_LARGEFILE);
		do_init();
		do_close("Init");
	} else
	if (argc == 3 && !strcmp(argv[1], "crash")) {
		do_open(argv[2], O_WRONLY|O_DIRECT|O_LARGEFILE|O_SYNC);
		do_crash(0);
		do_close("Crashing");
	} else
	if (argc == 3 && !strcmp(argv[1], "init-and-crash")) {
		do_open(argv[2], O_WRONLY|O_LARGEFILE);
		do_init();
		do_close("Init");
		do_open(argv[2], O_WRONLY|O_DIRECT|O_LARGEFILE|O_SYNC);
		do_crash(0);
		do_close("Crashing");
	} else
	if (argc == 3 && !strcmp(argv[1], "check")) {
		do_open(argv[2], O_RDONLY);
		do_check();
		do_close("Checking");
	} else
	if (argc == 3 && !strcmp(argv[1], "auto-noinit")) {
		goto auto_noinit;
	} else
	if (argc == 3 && !strcmp(argv[1], "auto")) {
		do_open(argv[2], O_WRONLY|O_LARGEFILE);
		do_init();
		do_close("Init");
auto_noinit:;
		int iter;
		bool finished = false;
		bool found_errors = false;
		for (iter=0; !finished; iter++) {
			do_open(argv[2], O_WRONLY|O_DIRECT|O_LARGEFILE|O_SYNC);
			if (do_crash(opt_autocycles * dev_size) == crash_status_gotctrlc)
				finished = true;
			do_close("Crashing");
			do_open(argv[2], O_RDONLY);
			if (do_check())
				finished = true, found_errors = true;
			do_close("Checking");
		}
		printf("Finished auto-mode after %d iterations.%s\n", iter, found_errors ? " FOUND ERRORS!" : "");
	} else {
print_usage:
		printf("Usage: %s [-p N] [-r N] [-c N] { init | crash | init-and-crash | auto | check } <dev>\n", progname);
		return 1;
	}
	return 0;
}