[disktools.git] / drivescan.c

/*
 * Copyright (C) 2004 Jens Axboe <axboe@suse.de>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 */
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <signal.h>
#include <string.h>
#include <getopt.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/ioctl.h>
#include <sys/time.h>
#include <sys/stat.h>

enum {
	arch_x86_64,
	arch_i386,
	arch_ppc,
	arch_ia64,
};

#if defined(__x86_64)
#define BITS_PER_LONG 64
static int arch = arch_x86_64;
inline unsigned long ffz(unsigned long bitmask)
{
	__asm__("bsfq %1,%0" :"=r" (bitmask) :"r" (~bitmask));
	return bitmask;
}
#elif defined (__i386__)
#define BITS_PER_LONG 32
static int arch = arch_i386;
inline unsigned long ffz(unsigned long bitmask)
{
	__asm__("bsfl %1,%0" :"=r" (bitmask) :"r" (~bitmask));
	return bitmask;
}
#elif defined (__powerpc__)
#define BITS_PER_LONG 32
static int arch = arch_ppc;
inline int __ilog2(unsigned long bitmask)
{
	int lz;

	asm ("cntlzw %0,%1" : "=r" (lz) : "r" (bitmask));
	return 31 - lz;
}

static inline int ffz(unsigned long bitmask)
{
	if ((bitmask = ~bitmask) == 0)
		return 32;
	return  __ilog2(bitmask & -bitmask);
}
#elif defined(__ia64__)
#define BITS_PER_LONG 64
static int arch = arch_ia64;
#define ia64_popcnt(x)							\
({									\
	unsigned long ia64_intri_res;					\
	asm ("popcnt %0=%1" : "=r" (ia64_intri_res) : "r" (x));		\
	ia64_intri_res;							\
})
inline unsigned long ffz(unsigned long bitmask)
{
	return ia64_popcnt(bitmask & (~bitmask - 1));
}
#else
#error "Unsupported architecture"
#endif

#if BITS_PER_LONG == 32
typedef unsigned long long u64;
typedef unsigned long u32;
#elif BITS_PER_LONG == 64
typedef unsigned long u64;
typedef unsigned int u32;
#endif

#ifndef BLKGETSIZE64
#define BLKGETSIZE64 _IOR(0x12,114,size_t)
#endif

#define BS	(65536)
#define MASK	(4095)

#define ALIGN(buf)	(((unsigned long) (buf) + MASK) & ~(MASK))

#define BITMAP_BLOCKS(blocks)	(((blocks) / BITS_PER_LONG) + 1)
#define BITMAP_SIZE(blocks)	(BITMAP_BLOCKS(blocks) * sizeof(unsigned long))

#define RANDOM_COVERAGE		(90)

static char device[256];
static char state_file[256];
static int progress_interval;
static int block_size;
static int restore_state = 1;
static struct timeval start;
static int sequential;
static int disk_util = 100;
static int disk_util_set;
static int state_file_fd = -1;
static int state_restored;

static unsigned long *read_bitmap;
static unsigned long find_next_bit_index;
static unsigned long blocks_read, blocks_read_now;
static u64 dev_size;

static struct drand48_data random_state;

struct our_time {
	unsigned long hour;
	unsigned long min;
	unsigned long sec;
};

#define DRIVESCAN_MAGIC		0xabc0ffee
#define DRIVESCAN_VERSION	2

struct progress_state {
	u32 magic;
	u32 arch;
	u32 version;
	u32 bs;
	u64 blocks;
	u32 disk_util;
	struct drand48_data random_state;
};

inline int test_block_read(unsigned long block)
{
	unsigned long index, bitnr;

	index = block / BITS_PER_LONG;
	bitnr = block & (BITS_PER_LONG - 1);
	return (read_bitmap[index] & (1UL << bitnr)) != 0;
}
	
inline int mark_block_read(unsigned long block)
{
        unsigned long index, bitnr, retval;

	index = block / BITS_PER_LONG;
	bitnr = block & (BITS_PER_LONG - 1);

	retval = (read_bitmap[index] & (1UL << bitnr)) != 0;
	read_bitmap[index] |= 1UL << bitnr;
	return retval;
}

inline unsigned long find_next_unread_block(void)
{
	unsigned int i, bitnr;

	for (i = find_next_bit_index; i < BITMAP_BLOCKS(dev_size); i++) {
		if (read_bitmap[i] == ~0UL)
			continue;

		goto gotit;
	}

	return -1;
gotit:
	find_next_bit_index = i;
	bitnr = ffz(read_bitmap[i]);
	read_bitmap[i] |= 1UL << bitnr;
	return bitnr + BITS_PER_LONG * i;
}

int load_progress_state(void)
{
	struct progress_state state;
	unsigned long i;
	int ret;

	if (state_file_fd == -1)
		return 1;

	if (lseek(state_file_fd, 0, SEEK_SET) == -1) {
		perror("state file seek");
		return 1;
	}

	ret = read(state_file_fd, &state, sizeof(state));
	if (ret == -1) {
		perror("load progress file");
		return 1;
	} else if (!ret)
		return 0;

	if (state.magic != DRIVESCAN_MAGIC) {
		printf("Bad magic %lx\n", (unsigned long) state.magic);
		return 1;
	}
	if (state.version != DRIVESCAN_VERSION) {
		printf("Unsupported file version (%d)\n", (int) state.version);
		return 1;
	}
	if (state.arch != arch) {
		printf("State file from different arch\n");
		return 1;
	}

	read_bitmap = malloc(BITMAP_SIZE(state.blocks));
	if (!read_bitmap) {
		printf("Unable to allocate bitmap (blocks=%Lu)\n", (unsigned long long) state.blocks);
		return 1;
	}
	memset(read_bitmap, 0, BITMAP_SIZE(dev_size));

	if (read(state_file_fd, (char *) read_bitmap, BITMAP_SIZE(state.blocks)) < BITMAP_SIZE(state.blocks)) {
		printf("Unable to read bitmap\n");
		return 1;
	}

	blocks_read = 0;
	for (i = 0; i < state.blocks; i++)
		if (test_block_read(i))
			blocks_read++;

	if (blocks_read == dev_size)
		return 0;

	memcpy(&random_state, &state.random_state, sizeof(random_state));

	dev_size = state.blocks;
	block_size = state.bs;
	if (!disk_util_set)
		disk_util = state.disk_util;
	printf("Restored state: %lu blocks\n", blocks_read);
	state_restored = 1;
	return 0;
}

void save_progress_state(void)
{
	struct progress_state state;
	int ret;

	if (ftruncate(state_file_fd, 0) == -1) {
		perror("truncate state file");
		return;
	}
	if (lseek(state_file_fd, 0, SEEK_SET) == -1) {
		perror("lseek state file");
		return;
	}

	memset(&state, 0, sizeof(state));
	state.magic = DRIVESCAN_MAGIC;
	state.arch = arch;
	state.version = DRIVESCAN_VERSION;
	state.bs = block_size;
	state.blocks = dev_size;
	state.disk_util = disk_util;
	memcpy(&state.random_state, &random_state, sizeof(random_state));

	ret = write(state_file_fd, &state, sizeof(state));
	if (ret == -1) {
		perror("write state file");
		return;
	}

	ret = write(state_file_fd, (char *) read_bitmap, BITMAP_SIZE(dev_size));
	if (ret == -1) {
		perror("write state bitmap");
		return;
	}

	fsync(state_file_fd);
	printf("Saved state: %lu blocks\n", blocks_read);
}

unsigned long utime_since(struct timeval *s)
{
	double sec, usec;
	struct timeval now;

	gettimeofday(&now, NULL);
	sec = now.tv_sec - s->tv_sec;
	usec = now.tv_usec - s->tv_usec;
	if (sec > 0 && usec < 0) {
		sec--;
		usec += 1000000;
	}
	return usec + (sec * (double) 1000000);
}

unsigned long time_since_now(void)
{
	double sec, usec, ret;
	struct timeval now;

	gettimeofday(&now, NULL);
	sec = now.tv_sec - start.tv_sec;
	usec = now.tv_usec - start.tv_usec;
	if (sec > 0 && usec < 0) {
		sec--;
		usec += 1000000;
	}
	ret = sec + usec / (double) 1000000;
	if (ret < 0)
		ret = 0;

	return (unsigned long) ret;
}

void normalize_time(struct our_time *ot, unsigned long seconds)
{
	ot->hour = seconds / 3600;
	seconds -= ot->hour * 3600;
	ot->min = seconds / 60;
	seconds -= ot->min * 60;
	ot->sec = seconds;
}

void show_progress_state(void)
{
	unsigned long rate, now;
	struct our_time comp, left;

	now = time_since_now();

	normalize_time(&comp, now);

	if (now == 0)
		now = 1;

	rate = blocks_read_now / now;

	memset(&left, 0, sizeof(left));
	if (blocks_read_now)
		normalize_time(&left, (dev_size * now / blocks_read_now) - now);
	
	printf("Status:\n");
	printf("\tRun: [%luh%lum%lusec, %luh%lum%lusec], %lu blocks/sec, %lu KiB/sec\n", comp.hour, comp.min, comp.sec, left.hour, left.min, left.sec, rate, rate * block_size / 1024);
	printf("\tRead:[%lu, %lu]: %lu%% coverage\n", blocks_read, (unsigned long) dev_size, (unsigned long) (blocks_read * 100 / dev_size));
}

void sig_handler(int sig)
{
	show_progress_state();

	switch (sig) {
		case SIGINT:
			if (state_file[0])
				save_progress_state();
			exit(0);
		case SIGALRM:
			if (progress_interval)
				alarm(progress_interval);
			break;
	}
}

unsigned long get_next_offset(unsigned long blocks)
{
	unsigned long b, retries = dev_size;

	if (!sequential) {
		retries = dev_size;
		do {
			long r;
			lrand48_r(&random_state, &r);
			b = (1+(double) (blocks-1) * r / (RAND_MAX+1.0));

			if (!mark_block_read(b))
				break;
			
		} while (--retries);

		if (!retries) {
			printf("Gave up finding new block\n");
			sequential = 1;
			goto seq;
		}
	} else {
seq:
		b = find_next_unread_block();
	}

	blocks_read++;
	blocks_read_now++;
	return b;
}

void idle_drive(struct timeval *s)
{
	unsigned long msleep, sleep_time;
	double ratio;

	if (disk_util == 100)
		return;

	msleep = utime_since(s);
	ratio = (double) 100 / disk_util;
	sleep_time = msleep * ratio - msleep;
	usleep(sleep_time);
}

void read_error_report(int fd, char *buffer, unsigned long long offset)
{
	unsigned long long bad_block = offset >> 9;
	int blocks = block_size / 512;
	int i, ret;

	if (lseek(fd, offset, SEEK_SET) == -1) {
		perror("lseek");
		return;
	}

	for (i = 0; i < blocks; i++) {
		ret = read(fd, buffer, 512);

		if (ret == 512)
			continue;
		else if (ret < 512 && ret >= 0) {
			printf("Short read at sector %Lu\n", bad_block + i);
			continue;
		}
		if (errno == EINVAL) {
			printf("512-b O_DIRECT does not work\n");
			break;
		}
		if (errno == EIO)
			printf("Sector %Lu bad\n", bad_block + i);
	}
}

void do_random_reads(int fd, unsigned long blocks)
{
	struct timeval s;
	char *ptr, *buffer;
	int ret;

	ptr = malloc(block_size + MASK + 1);
	buffer = (char *) ALIGN(ptr);

	do {
		off_t offset = get_next_offset(blocks);
		int coverage;

		offset = get_next_offset(blocks);
		if (offset == -1UL)
			break;

		offset = offset * block_size;

		if (lseek(fd, offset, SEEK_SET) == -1) {
			perror("lseek");
			break;
		}

		gettimeofday(&s, NULL);
		ret = read(fd, buffer, block_size);

		if (ret == -1) {
			if (errno == EINVAL) {
				printf("EINVAL on read, check if O_DIRECT works\n");
				break;
			} else if (errno == EIO) {
				read_error_report(fd, buffer, offset);
			} else {
				perror("read");
				break;
			}
		}
		if (ret < block_size)
			read_error_report(fd, buffer, offset);

		coverage = blocks_read * 100 / dev_size;
		if (!sequential && coverage >= RANDOM_COVERAGE) {
			printf("Switching to scan for rest of blocks\n");
			sequential = 1;
		}

		if (blocks_read == dev_size)
			break;

		idle_drive(&s);

	} while (1);

	show_progress_state();

	free(ptr);
}

int seed_randomizer(void)
{
	unsigned long seed;
	int fd;

	fd = open("/dev/random", O_RDONLY);
	if (fd == -1) {
		perror("open rand");
		return 1;
	}

	if (read(fd, &seed, sizeof(seed)) < (int) sizeof(seed)) {
		perror("read rand");
		return 1;
	}

	close(fd);
	
	printf("Randomizer seed: %lx\n", seed);
	srand48_r(seed, &random_state);
	return 0;
}

void device_to_state_file(void)
{
	unsigned int i, offset;

	offset = sprintf(state_file, "scanstate");

	for (i = 0; i < strlen(device); i++) {
		if (device[i] == '/')
			state_file[offset+i] = '_';
		else
			state_file[offset+i] = device[i];
	}
}

void init_state_file(void)
{
	if (!state_file[0])
		device_to_state_file();

	state_file_fd = open(state_file, O_RDWR | O_CREAT, 0644);

	if (state_file_fd == -1) {
		perror("open state file");
		return;
	}
}

int get_options(int argc, char *argv[])
{
	static struct option longoptions[] = {
		{ "device",	1, 	NULL,	'd' },
		{ "block-size",	1,	NULL,	'b' },
		{ "progress",	1,	NULL,	'p' },
		{ "file",	1,	NULL,	'f' },
		{ "restore",	1,	NULL,	'r' },
		{ "utilization",1,	NULL,	'i' },
	};
	int c, res = 0;

	while ((c = getopt_long(argc, argv, "d:b:p:f:ru:", longoptions, &res)) != -1) {
		switch (c) {
			case 'd':
				strncpy(device, optarg, sizeof(device) - 1);
				break;
			case 'b':
				block_size = atoi(optarg);
				if (block_size < 1 || block_size > 1024) {
					printf("bad block size %d\n", block_size);
					block_size = BS;
				}
				block_size <<= 10;
				break;
			case 'p':
				progress_interval = atoi(optarg);
				break;
			case 'f':
				strncpy(state_file, optarg, sizeof(state_file)-1);
				break;
			case 'r':
				restore_state = !!atoi(optarg);
				break;
			case 'u':
				disk_util = atoi(optarg);
				if (disk_util < 0)
					disk_util = 1;
				if (disk_util > 100)
					disk_util = 100;
				disk_util_set = 1;
				break;
		}
	}

	if (!device[0]) {
		printf("Must give device with -d\n");
		return 1;
	}
	if (!block_size)
		block_size = BS;

	return 0;
}

int main(int argc, char *argv[])
{
	u64 blocks, size;
	struct stat statbuf;
	int fd;

	if (get_options(argc, argv))
		return 1;

	seed_randomizer();

	init_state_file();

	if (restore_state)
		load_progress_state();

	if (stat(device, &statbuf) == -1) {
		perror("stat");
		return 1;
	}

	if (!S_ISBLK(statbuf.st_mode)) {
		printf("%s does not appear to be a block device\n", device);
		return 1;
	}

	fd = open(device, O_RDONLY | O_DIRECT | O_LARGEFILE);
	if (fd == -1) {
		perror("open");
		return 2;
	}

	if (ioctl(fd, BLKGETSIZE64, &size) == -1) {
		perror("BLKGETSIZE64");
		return 3;
	}

	blocks = size / block_size;

	if (state_restored && dev_size != blocks) {
		printf("Device size seems different %lu != %lu\n", (unsigned long) dev_size, (unsigned long) blocks);
		return 1;
	}

	dev_size = blocks;

	if (!read_bitmap) {
		read_bitmap = malloc(BITMAP_SIZE(dev_size));
		memset(read_bitmap, 0, BITMAP_SIZE(dev_size));
	}

	signal(SIGUSR1, sig_handler);
	signal(SIGINT, sig_handler);

	if (progress_interval) {
		signal(SIGALRM, sig_handler);
		alarm(progress_interval);
	}

	gettimeofday(&start, NULL);
	do_random_reads(fd, blocks);

	save_progress_state();

	free((void *) read_bitmap);
	close(fd);
	if (state_file_fd != -1)
		close(state_file_fd);

	return 0;
}
Commit	Line	Data
854a84cf JA	1	/*
	2	* Copyright (C) 2004 Jens Axboe <axboe@suse.de>
	3	*
	4	* This program is free software; you can redistribute it and/or modify
	5	* it under the terms of the GNU General Public License version 2 as
	6	* published by the Free Software Foundation.
	7	*/
	8	#include <stdio.h>
	9	#include <stdlib.h>
	10	#include <fcntl.h>
	11	#include <unistd.h>
	12	#include <signal.h>
	13	#include <string.h>
	14	#include <getopt.h>
	15	#include <errno.h>
	16	#include <sys/types.h>
	17	#include <sys/ioctl.h>
	18	#include <sys/time.h>
	19	#include <sys/stat.h>
	20
	21	enum {
	22	arch_x86_64,
	23	arch_i386,
	24	arch_ppc,
	25	arch_ia64,
	26	};
	27
	28	#if defined(__x86_64)
	29	#define BITS_PER_LONG 64
	30	static int arch = arch_x86_64;
	31	inline unsigned long ffz(unsigned long bitmask)
	32	{
	33	__asm__("bsfq %1,%0" :"=r" (bitmask) :"r" (~bitmask));
	34	return bitmask;
	35	}
	36	#elif defined (__i386__)
	37	#define BITS_PER_LONG 32
	38	static int arch = arch_i386;
	39	inline unsigned long ffz(unsigned long bitmask)
	40	{
	41	__asm__("bsfl %1,%0" :"=r" (bitmask) :"r" (~bitmask));
	42	return bitmask;
	43	}
	44	#elif defined (__powerpc__)
	45	#define BITS_PER_LONG 32
	46	static int arch = arch_ppc;
	47	inline int __ilog2(unsigned long bitmask)
	48	{
	49	int lz;
	50
	51	asm ("cntlzw %0,%1" : "=r" (lz) : "r" (bitmask));
	52	return 31 - lz;
	53	}
	54
	55	static inline int ffz(unsigned long bitmask)
	56	{
	57	if ((bitmask = ~bitmask) == 0)
	58	return 32;
	59	return __ilog2(bitmask & -bitmask);
	60	}
	61	#elif defined(__ia64__)
	62	#define BITS_PER_LONG 64
	63	static int arch = arch_ia64;
	64	#define ia64_popcnt(x) \
65	({ \
66	unsigned long ia64_intri_res; \
67	asm ("popcnt %0=%1" : "=r" (ia64_intri_res) : "r" (x)); \
68	ia64_intri_res; \
69	})
70	inline unsigned long ffz(unsigned long bitmask)
71	{
72	return ia64_popcnt(bitmask & (~bitmask - 1));
73	}
74	#else
75	#error "Unsupported architecture"
76	#endif
77
78	#if BITS_PER_LONG == 32
79	typedef unsigned long long u64;
80	typedef unsigned long u32;
81	#elif BITS_PER_LONG == 64
82	typedef unsigned long u64;
83	typedef unsigned int u32;
84	#endif
85
86	#ifndef BLKGETSIZE64
87	#define BLKGETSIZE64 _IOR(0x12,114,size_t)
88	#endif
89
f61c31f7	90	#define BS (65536)
854a84cf JA	91	#define MASK (4095)
	92
	93	#define ALIGN(buf) (((unsigned long) (buf) + MASK) & ~(MASK))
	94
	95	#define BITMAP_BLOCKS(blocks) (((blocks) / BITS_PER_LONG) + 1)
	96	#define BITMAP_SIZE(blocks) (BITMAP_BLOCKS(blocks) * sizeof(unsigned long))
	97
	98	#define RANDOM_COVERAGE (90)
	99
	100	static char device[256];
	101	static char state_file[256];
	102	static int progress_interval;
	103	static int block_size;
	104	static int restore_state = 1;
	105	static struct timeval start;
	106	static int sequential;
	107	static int disk_util = 100;
c040c171	108	static int disk_util_set;
854a84cf JA	109	static int state_file_fd = -1;
	110	static int state_restored;
	111
	112	static unsigned long *read_bitmap;
	113	static unsigned long find_next_bit_index;
	114	static unsigned long blocks_read, blocks_read_now;
	115	static u64 dev_size;
	116
	117	static struct drand48_data random_state;
	118
	119	struct our_time {
	120	unsigned long hour;
	121	unsigned long min;
	122	unsigned long sec;
	123	};
	124
	125	#define DRIVESCAN_MAGIC 0xabc0ffee
	126	#define DRIVESCAN_VERSION 2
	127
	128	struct progress_state {
	129	u32 magic;
	130	u32 arch;
	131	u32 version;
	132	u32 bs;
	133	u64 blocks;
	134	u32 disk_util;
	135	struct drand48_data random_state;
	136	};
	137
	138	inline int test_block_read(unsigned long block)
	139	{
	140	unsigned long index, bitnr;
	141
	142	index = block / BITS_PER_LONG;
	143	bitnr = block & (BITS_PER_LONG - 1);
	144	return (read_bitmap[index] & (1UL << bitnr)) != 0;
	145	}
	146
	147	inline int mark_block_read(unsigned long block)
	148	{
	149	unsigned long index, bitnr, retval;
	150
	151	index = block / BITS_PER_LONG;
	152	bitnr = block & (BITS_PER_LONG - 1);
	153
	154	retval = (read_bitmap[index] & (1UL << bitnr)) != 0;
	155	read_bitmap[index] \|= 1UL << bitnr;
	156	return retval;
	157	}
	158
	159	inline unsigned long find_next_unread_block(void)
	160	{
	161	unsigned int i, bitnr;
	162
	163	for (i = find_next_bit_index; i < BITMAP_BLOCKS(dev_size); i++) {
	164	if (read_bitmap[i] == ~0UL)
	165	continue;
	166
	167	goto gotit;
	168	}
	169
	170	return -1;
	171	gotit:
	172	find_next_bit_index = i;
173	bitnr = ffz(read_bitmap[i]);
174	read_bitmap[i] \|= 1UL << bitnr;
175	return bitnr + BITS_PER_LONG * i;
176	}
177
178	int load_progress_state(void)
179	{
180	struct progress_state state;
181	unsigned long i;
182	int ret;
183
184	if (state_file_fd == -1)
185	return 1;
186
187	if (lseek(state_file_fd, 0, SEEK_SET) == -1) {
188	perror("state file seek");
189	return 1;
190	}
191
192	ret = read(state_file_fd, &state, sizeof(state));
193	if (ret == -1) {
194	perror("load progress file");
195	return 1;
196	} else if (!ret)
197	return 0;
198
199	if (state.magic != DRIVESCAN_MAGIC) {
200	printf("Bad magic %lx\n", (unsigned long) state.magic);
201	return 1;
202	}
203	if (state.version != DRIVESCAN_VERSION) {
204	printf("Unsupported file version (%d)\n", (int) state.version);
205	return 1;
206	}
207	if (state.arch != arch) {
208	printf("State file from different arch\n");
209	return 1;
210	}
211
212	read_bitmap = malloc(BITMAP_SIZE(state.blocks));
213	if (!read_bitmap) {
214	printf("Unable to allocate bitmap (blocks=%Lu)\n", (unsigned long long) state.blocks);
215	return 1;
216	}
217	memset(read_bitmap, 0, BITMAP_SIZE(dev_size));
218
219	if (read(state_file_fd, (char *) read_bitmap, BITMAP_SIZE(state.blocks)) < BITMAP_SIZE(state.blocks)) {
220	printf("Unable to read bitmap\n");
221	return 1;
222	}
223
224	blocks_read = 0;
225	for (i = 0; i < state.blocks; i++)
226	if (test_block_read(i))
227	blocks_read++;
228
229	if (blocks_read == dev_size)
230	return 0;
231
232	memcpy(&random_state, &state.random_state, sizeof(random_state));
233
234	dev_size = state.blocks;
235	block_size = state.bs;
c040c171 JA	236	if (!disk_util_set)
c040c171 JA	237	disk_util = state.disk_util;
854a84cf JA	238	printf("Restored state: %lu blocks\n", blocks_read);
	239	state_restored = 1;
	240	return 0;
	241	}
	242
	243	void save_progress_state(void)
	244	{
	245	struct progress_state state;
	246	int ret;
	247
	248	if (ftruncate(state_file_fd, 0) == -1) {
	249	perror("truncate state file");
	250	return;
	251	}
	252	if (lseek(state_file_fd, 0, SEEK_SET) == -1) {
	253	perror("lseek state file");
	254	return;
	255	}
	256
	257	memset(&state, 0, sizeof(state));
	258	state.magic = DRIVESCAN_MAGIC;
	259	state.arch = arch;
	260	state.version = DRIVESCAN_VERSION;
	261	state.bs = block_size;
	262	state.blocks = dev_size;
	263	state.disk_util = disk_util;
	264	memcpy(&state.random_state, &random_state, sizeof(random_state));
	265
	266	ret = write(state_file_fd, &state, sizeof(state));
	267	if (ret == -1) {
	268	perror("write state file");
	269	return;
	270	}
	271
	272	ret = write(state_file_fd, (char *) read_bitmap, BITMAP_SIZE(dev_size));
	273	if (ret == -1) {
	274	perror("write state bitmap");
	275	return;
	276	}
	277
	278	fsync(state_file_fd);
	279	printf("Saved state: %lu blocks\n", blocks_read);
	280	}
	281
	282	unsigned long utime_since(struct timeval *s)
	283	{
	284	double sec, usec;
	285	struct timeval now;
	286
	287	gettimeofday(&now, NULL);
	288	sec = now.tv_sec - s->tv_sec;
	289	usec = now.tv_usec - s->tv_usec;
	290	if (sec > 0 && usec < 0) {
	291	sec--;
	292	usec += 1000000;
	293	}
	294	return usec + (sec * (double) 1000000);
	295	}
	296
	297	unsigned long time_since_now(void)
	298	{
	299	double sec, usec, ret;
	300	struct timeval now;
	301
302	gettimeofday(&now, NULL);
303	sec = now.tv_sec - start.tv_sec;
304	usec = now.tv_usec - start.tv_usec;
305	if (sec > 0 && usec < 0) {
306	sec--;
307	usec += 1000000;
308	}
309	ret = sec + usec / (double) 1000000;
310	if (ret < 0)
311	ret = 0;
312
313	return (unsigned long) ret;
314	}
315
316	void normalize_time(struct our_time *ot, unsigned long seconds)
317	{
318	ot->hour = seconds / 3600;
319	seconds -= ot->hour * 3600;
320	ot->min = seconds / 60;
321	seconds -= ot->min * 60;
322	ot->sec = seconds;
323	}
324
325	void show_progress_state(void)
326	{
327	unsigned long rate, now;
328	struct our_time comp, left;
329
330	now = time_since_now();
331
332	normalize_time(&comp, now);
333
334	if (now == 0)
335	now = 1;
336
337	rate = blocks_read_now / now;
338
339	memset(&left, 0, sizeof(left));
340	if (blocks_read_now)
341	normalize_time(&left, (dev_size * now / blocks_read_now) - now);
342
343	printf("Status:\n");
344	printf("\tRun: [%luh%lum%lusec, %luh%lum%lusec], %lu blocks/sec, %lu KiB/sec\n", comp.hour, comp.min, comp.sec, left.hour, left.min, left.sec, rate, rate * block_size / 1024);
345	printf("\tRead:[%lu, %lu]: %lu%% coverage\n", blocks_read, (unsigned long) dev_size, (unsigned long) (blocks_read * 100 / dev_size));
346	}
347
348	void sig_handler(int sig)
349	{
350	show_progress_state();
351
352	switch (sig) {
353	case SIGINT:
354	if (state_file[0])
355	save_progress_state();
356	exit(0);
357	case SIGALRM:
358	if (progress_interval)
359	alarm(progress_interval);
360	break;
361	}
362	}
363
364	unsigned long get_next_offset(unsigned long blocks)
365	{
366	unsigned long b, retries = dev_size;
367
368	if (!sequential) {
369	retries = dev_size;
370	do {
371	long r;
372	lrand48_r(&random_state, &r);
373	b = (1+(double) (blocks-1) * r / (RAND_MAX+1.0));
374
375	if (!mark_block_read(b))
376	break;
377
378	} while (--retries);
379
380	if (!retries) {
381	printf("Gave up finding new block\n");
382	sequential = 1;
383	goto seq;
384	}
385	} else {
386	seq:
387	b = find_next_unread_block();
388	}
389
390	blocks_read++;
391	blocks_read_now++;
392	return b;
393	}
394
395	void idle_drive(struct timeval *s)
396	{
397	unsigned long msleep, sleep_time;
398	double ratio;
399
400	if (disk_util == 100)
401	return;
402
403	msleep = utime_since(s);
404	ratio = (double) 100 / disk_util;
405	sleep_time = msleep * ratio - msleep;
406	usleep(sleep_time);
407	}
408
409	void read_error_report(int fd, char *buffer, unsigned long long offset)
410	{
411	unsigned long long bad_block = offset >> 9;
412	int blocks = block_size / 512;
413	int i, ret;
414
415	if (lseek(fd, offset, SEEK_SET) == -1) {
416	perror("lseek");
417	return;
418	}
419
420	for (i = 0; i < blocks; i++) {
421	ret = read(fd, buffer, 512);
422
423	if (ret == 512)
424	continue;
425	else if (ret < 512 && ret >= 0) {
426	printf("Short read at sector %Lu\n", bad_block + i);
427	continue;
428	}
429	if (errno == EINVAL) {
430	printf("512-b O_DIRECT does not work\n");
431	break;
432	}
433	if (errno == EIO)
434	printf("Sector %Lu bad\n", bad_block + i);
435	}
436	}
437
438	void do_random_reads(int fd, unsigned long blocks)
439	{
440	struct timeval s;
441	char ptr, buffer;
442	int ret;
443
444	ptr = malloc(block_size + MASK + 1);
445	buffer = (char *) ALIGN(ptr);
446
447	do {
448	off_t offset = get_next_offset(blocks);
449	int coverage;
450
451	offset = get_next_offset(blocks);
452	if (offset == -1UL)
453	break;
454
455	offset = offset * block_size;
456
457	if (lseek(fd, offset, SEEK_SET) == -1) {
458	perror("lseek");
459	break;
460	}
461
462	gettimeofday(&s, NULL);
463	ret = read(fd, buffer, block_size);
464
465	if (ret == -1) {
466	if (errno == EINVAL) {
467	printf("EINVAL on read, check if O_DIRECT works\n");
468	break;
469	} else if (errno == EIO) {
470	read_error_report(fd, buffer, offset);
471	} else {
472	perror("read");
473	break;
474	}
475	}
476	if (ret < block_size)
477	read_error_report(fd, buffer, offset);
478
479	coverage = blocks_read * 100 / dev_size;
480	if (!sequential && coverage >= RANDOM_COVERAGE) {
481	printf("Switching to scan for rest of blocks\n");
482	sequential = 1;
483	}
484
485	if (blocks_read == dev_size)
486	break;
487
488	idle_drive(&s);
489
490	} while (1);
491
492	show_progress_state();
493
494	free(ptr);
495	}
496
497	int seed_randomizer(void)
498	{
499	unsigned long seed;
500	int fd;
501
502	fd = open("/dev/random", O_RDONLY);
503	if (fd == -1) {
504	perror("open rand");
505	return 1;
506	}
507
508	if (read(fd, &seed, sizeof(seed)) < (int) sizeof(seed)) {
509	perror("read rand");
510	return 1;
511	}
512
513	close(fd);
514
515	printf("Randomizer seed: %lx\n", seed);
516	srand48_r(seed, &random_state);
517	return 0;
518	}
519
520	void device_to_state_file(void)
521	{
522	unsigned int i, offset;
523
524	offset = sprintf(state_file, "scanstate");
525
526	for (i = 0; i < strlen(device); i++) {
527	if (device[i] == '/')
528	state_file[offset+i] = '_';
529	else
530	state_file[offset+i] = device[i];
531	}
532	}
533
534	void init_state_file(void)
535	{
536	if (!state_file[0])
537	device_to_state_file();
538
539	state_file_fd = open(state_file, O_RDWR \| O_CREAT, 0644);
540
541	if (state_file_fd == -1) {
542	perror("open state file");
543	return;
544	}
545	}
546
547	int get_options(int argc, char *argv[])
548	{
549	static struct option longoptions[] = {
550	{ "device", 1, NULL, 'd' },
254abedc	551	{ "block-size", 1, NULL, 'b' },
854a84cf JA	552	{ "progress", 1, NULL, 'p' },
	553	{ "file", 1, NULL, 'f' },
	554	{ "restore", 1, NULL, 'r' },
	555	{ "utilization",1, NULL, 'i' },
	556	};
	557	int c, res = 0;
	558
	559	while ((c = getopt_long(argc, argv, "d:b:p:f:ru:", longoptions, &res)) != -1) {
	560	switch (c) {
	561	case 'd':
	562	strncpy(device, optarg, sizeof(device) - 1);
	563	break;
	564	case 'b':
	565	block_size = atoi(optarg);
	566	if (block_size < 1 \|\| block_size > 1024) {
	567	printf("bad block size %d\n", block_size);
	568	block_size = BS;
	569	}
	570	block_size <<= 10;
	571	break;
	572	case 'p':
	573	progress_interval = atoi(optarg);
	574	break;
	575	case 'f':
	576	strncpy(state_file, optarg, sizeof(state_file)-1);
	577	break;
	578	case 'r':
	579	restore_state = !!atoi(optarg);
	580	break;
	581	case 'u':
	582	disk_util = atoi(optarg);
	583	if (disk_util < 0)
	584	disk_util = 1;
	585	if (disk_util > 100)
	586	disk_util = 100;
c040c171	587	disk_util_set = 1;
854a84cf JA	588	break;
	589	}
	590	}
	591
	592	if (!device[0]) {
	593	printf("Must give device with -d\n");
	594	return 1;
	595	}
	596	if (!block_size)
	597	block_size = BS;
	598
	599	return 0;
	600	}
	601
	602	int main(int argc, char *argv[])
	603	{
	604	u64 blocks, size;
	605	struct stat statbuf;
	606	int fd;
	607
	608	if (get_options(argc, argv))
	609	return 1;
	610
	611	seed_randomizer();
	612
	613	init_state_file();
	614
	615	if (restore_state)
	616	load_progress_state();
	617
	618	if (stat(device, &statbuf) == -1) {
	619	perror("stat");
	620	return 1;
	621	}
	622
	623	if (!S_ISBLK(statbuf.st_mode)) {
	624	printf("%s does not appear to be a block device\n", device);
	625	return 1;
	626	}
	627
	628	fd = open(device, O_RDONLY \| O_DIRECT \| O_LARGEFILE);
	629	if (fd == -1) {
	630	perror("open");
	631	return 2;
	632	}
	633
	634	if (ioctl(fd, BLKGETSIZE64, &size) == -1) {
	635	perror("BLKGETSIZE64");
	636	return 3;
	637	}
	638
	639	blocks = size / block_size;
	640
	641	if (state_restored && dev_size != blocks) {
	642	printf("Device size seems different %lu != %lu\n", (unsigned long) dev_size, (unsigned long) blocks);
	643	return 1;
	644	}
	645
	646	dev_size = blocks;
	647
	648	if (!read_bitmap) {
	649	read_bitmap = malloc(BITMAP_SIZE(dev_size));
	650	memset(read_bitmap, 0, BITMAP_SIZE(dev_size));
	651	}
652
653	signal(SIGUSR1, sig_handler);
654	signal(SIGINT, sig_handler);
655
656	if (progress_interval) {
657	signal(SIGALRM, sig_handler);
658	alarm(progress_interval);
659	}
660
661	gettimeofday(&start, NULL);
662	do_random_reads(fd, blocks);
663
664	save_progress_state();
665
666	free((void *) read_bitmap);
667	close(fd);
668	if (state_file_fd != -1)
669	close(state_file_fd);
670
671	return 0;
672	}