#!/usr/bin/perl
#
# Brad's el-ghetto do-our-storage-stacks-lie?-script
#
# https://gist.github.com/bradfitz/3172656#file-diskchecker-pl
 
sub usage {
	die <<'END';
	Usage: diskchecker.pl -s <server[:port]> verify <file>
	diskchecker.pl -s <server[:port]> create <file> <size_in_MB>
	diskchecker.pl -l [port]
	END
}
 
use strict;
use IO::Socket::INET;
use IO::Handle;
use Getopt::Long;
use Socket qw(IPPROTO_TCP TCP_NODELAY);
 
my $server;
my $listen;
usage() unless GetOptions('server=s' => \$server,
	'listen:5400' => \$listen);
usage() unless $server || $listen;
usage() if     $server && $listen;
 
# LISTEN MODE:
listen_mode($listen) if $listen;
 
# CLIENT MODE:
my $LEN = 16 * 1024;   # 16kB (same as InnoDB page)
my $mode = shift;
usage() unless $mode =~ /^verify|create$/;
 
my $file = shift or usage();
my $size;
if ($mode eq "create") {
	$size = shift or usage();
}
 
$server .= ":5400" unless $server =~ /:/;
 
my $sock = IO::Socket::INET->new(PeerAddr => $server)
	or die "Couldn't connect to host:port of '$server'\n";
 
setsockopt($sock, IPPROTO_TCP, TCP_NODELAY, pack("l", 1)) or die;
 
create() if $mode eq "create";
verify() if $mode eq "verify";
exit 0;
 
sub verify {
	sendmsg($sock, "read");
 
	my $error_ct = 0;
	my %error_ct;
 
	my $size = -s $file;
	my $max_pages = int($size / $LEN);
 
	my $percent;
	my $last_dump = 0;
	my $show_percent = sub {
		printf " verifying: %.02f%%\n", $percent;
	};
 
	open (F, $file) or die "Couldn't open file $file for read\n";
 
	while (<$sock>) {
		chomp;
		my ($page, $good, $val, $ago) = split(/\t/, $_);
		$percent = 100 * $page / ($max_pages || 1);
 
		my $now = time;
		if ($last_dump != $now) {
			$last_dump = $now;
			$show_percent->();
		}
 
		next unless $good;
 
		my $offset = $page * $LEN;
		sysseek F, $offset, 0;
		my $buf;
		my $rv = sysread(F, $buf, $LEN);
		my $tobe = sprintf("%08x", $val) x ($LEN / 8);
		substr($tobe, $LEN-1, 1) = "\n";
 
		unless ($buf eq $tobe) {
			$error_ct{$ago}++;
			$error_ct++;
			print "  Error at page $page, $ago seconds before end.\n";
		}
	}
	$show_percent->();
 
	print "Total errors: $error_ct\n";
	if ($error_ct) {
		print "Histogram of seconds before end:\n";
		foreach (sort { $a <=> $b } keys %error_ct) {
			printf "  %4d %4d\n", $_, $error_ct{$_};
		}
	}
}
 
sub create {
	open (F, ">$file") or die "Couldn't open file $file\n";
 
	my $ioh = IO::Handle->new_from_fd(fileno(F), "w")
		or die;
 
	my $pages = int( ($size * 1024 * 1024) / $LEN );  # 50 MiB of 16k pages (3200 pages)
 
	my %page_hit;
	my $pages_hit = 0;
	my $uniq_pages_hit = 0;
	my $start = time();
	my $last_dump = $start;
 
	while (1) {
		my $rand = int rand 2000000;
		my $buf = sprintf("%08x", $rand) x ($LEN / 8);
		substr($buf, $LEN-1, 1) = "\n";
 
		my $pagenum = int rand $pages;
		my $offset = $pagenum * $LEN;
 
		sendmsg($sock, "pre\t$pagenum\t$rand");
 
# now wait for acknowledgement
		my $ok = readmsg($sock);
		die "didn't get 'ok' from server ($pagenum $rand), msg=[$ok] = $!" unless $ok eq "ok";
 
		sysseek F,$offset,0;
		my $wv = syswrite(F, $buf, $LEN);
		die "return value wasn't $LEN\n" unless $wv == $LEN;
		$ioh->sync or die "couldn't do IO::Handle::sync";  # does fsync
 
		sendmsg($sock, "post\t$pagenum\t$rand");
 
		$pages_hit++;
		unless ($page_hit{$pagenum}++) {
			$uniq_pages_hit++;
		}
 
		my $now = time;
		if ($now != $last_dump) {
			$last_dump = $now;
			my $runtime = $now - $start;
			printf("  diskchecker: running %d sec, %.02f%% coverage of %d MB (%d writes; %d/s)\n",
				$runtime,
				(100 * $uniq_pages_hit / $pages),
				$size,
				$pages_hit,
				$pages_hit / $runtime,
			);
		}
	}
}
 
sub readmsg {
	my $sock = shift;
	my $len;
	my $rv = sysread($sock, $len, 1);
	return undef unless $rv == 1;
	my $msg;
	$rv = sysread($sock, $msg, ord($len));
	return $msg;
}
 
sub sendmsg {
	my ($sock, $msg) = @_;
	my $rv = syswrite($sock, chr(length($msg)) . $msg);
	my $expect = length($msg) + 1;
	die "sendmsg failed rv=$rv, expect=$expect" unless $rv == $expect;
	return 1;
}
 
sub listen_mode {
	my $port = shift;
	my $server = IO::Socket::INET->new(ReuseAddr => 1,
		Listen => 1,
		LocalPort => $port)
		or die "couldn't make server socket\n";
 
	while (1) {
		print "[server] diskchecker.pl: waiting for connection...\n";
		my $sock = $server->accept()
			or die "  die: no connection?";
		setsockopt($sock, IPPROTO_TCP, TCP_NODELAY, pack("l", 1)) or die;
 
		fork and next;
		process_incoming_conn($sock);
		exit 0;
	}
}
 
sub process_incoming_conn {
	my $sock = shift;
	my $peername = getpeername($sock) or
	die "connection not there?\n";
	my ($port, $iaddr) = sockaddr_in($peername);
	my $ip = inet_ntoa($iaddr);
 
	my $file = "/tmp/$ip.diskchecker";
	die "[$ip] $file is a symlink" if -l $file;
 
	print "[$ip] New connection\n";
 
	my $lines = 0;
	my %state;
	my $end;
 
	while (1) {
		if ($lines) {
			last unless wait_for_readability(fileno($sock), 3);
		}
		my $line = readmsg($sock);
		last unless $line;
 
		if ($line eq "read") {
			print "[$ip] Sending state info from ${ip}'s last create.\n";
			open (S, "$file") or die "Couldn't open $file for reading.";
			while (<S>) {
				print $sock $_;
			}
			close S;
			print "[$ip] Done.\n";
			exit 0;
		}
 
		$lines++;
		my $now = time;
		$end = $now;
		my ($state, $pagenum, $rand) = split(/\t/, $line);
		if ($state eq "pre") {
			$state{$pagenum} = [ 0, $rand+0, $now ];
			sendmsg($sock, "ok");
		} elsif ($state eq "post") {
			$state{$pagenum} = [ 1, $rand+0, $now ];
		}
		print "[$ip] $lines writes\n" if $lines % 1000 == 0;
	}
 
	print "[$ip] Writing state file...\n";
	open (S, ">$file") or die "Couldn't open $file for writing.";
	foreach (sort { $a <=> $b } keys %state) {
		my $v = $state{$_};
		my $before_end = $end - $v->[2];
		print S "$_\t$v->[0]\t$v->[1]\t$before_end\n";
	}
	print "[$ip] Done.\n";
}
 
sub wait_for_readability {
	my ($fileno, $timeout) = @_;
	return 0 unless $fileno && $timeout;
 
	my $rin;
	vec($rin, $fileno, 1) = 1;
	my $nfound = select($rin, undef, undef, $timeout);
 
	return 0 unless defined $nfound;
	return $nfound ? 1 : 0;
}