#!/usr/bin/perl
#
# PostgreSQL IO Monitor
# by Jeff Trout <jeff@jefftrout.com>
# http://postgresql.jefftrout.com/
#
# Provides a "top" style listing of what tables are not, what are not.
#
# Released under the BSD Public License
#
# Requires enable_stat_block enabled in your db
#
# TODO: look at specific index usage to identify a hot index (rather
#       than just an index on a table. (could we do w/ union?)
use DBI;
use POSIX;
$|=1;

$version = "0.1";

# config
$ioOnly = 0;        # do not count cache hits
$interval = 2;      # seconds
$dispInBlocks = undef;  # display output in blocks rather than kB
$dbStr = undef;     # db conn str
$dbUser = undef;    # db username
$dbPass = undef;    # db pass
$sortBy = "total";  # key to sort display by


$ARGC = scalar(@ARGV);

# process args
# don't forget to update usage() if you change it
#
for($curArg = 0; $curArg < $ARGC; $curArg++)
  {
	$arg = $ARGV[$curArg];

	if($arg eq "-v")
	  {
		print "PG IO Monitor v$version\n";
		exit;
	  }
	elsif($arg eq "-b")
	  {
		$dispInBlocks = 1;
	  }
	elsif($arg eq "-c")
	  {
		$ioOnly = 1;
		next;
	  }
	elsif($arg eq "-u")
	  {
		$dbUser = $ARGV[$curArg + 1];
		$curArg++;
		next;
	  }
	elsif($arg eq "-p")
	  {
		$dbPass = $ARGV[$curArg + 1];
		$curArg++;
		next;
	  }
	elsif($arg eq "-e")
	  {
		$sortBy = "heap";
	  }
	elsif($arg eq "-i")
	  {
		$sortBy = "index";
	  }
	elsif($arg eq "-t")
	  {
		$sortBy = "total";
	  }
	elsif($arg eq "-o")
	  {
		$sortBy = "toast";
	  }
	elsif($arg eq "-a")
	  {
		$sortBy = "toastidx";
	  }
	elsif($arg eq "-?")
	  {
		usage();
		exit;
	  }
	else
	  {
		$dbStr = $arg;
	  }
  }


# connect
$db = DBI->connect("dbi:Pg:".$dbStr, $dbUser, $dbPass);
if(!$db)
  {
	print "Cannot connect to the db: $dbstr\n";
	usage();
	exit;
  }

setupTerm();
$SIG{INT} = \&cleanupAndExit;
$SIG{TERM} = \&cleanupAndExit;

# get some info
$dbname = $db->{Name};

# Check to see if we have block stats enabled
$blockOn = @{$db->selectcol_arrayref("show stats_block_level")}->[0];
if($blockOn ne "on")
  {
	print "ERROR: you need to enable stats_block_level in postgresql.conf\n";
	exit;
  }

# get PG blocksize
$blockSize = @{$db->selectcol_arrayref("show block_size")}->[0];
$blockSize = $blockSize / 1024; # normalize to kB

# If we're not showing in block's get block size so we can show kB/sec
if($dispInBlocks)
  {
	$blockSizeModifier = 1;
	$unitDisp = "blk/sec";
  }
else
  {
	$blockSizeModifier = $blockSize;
	$unitDisp = "kB/sec";
  }


# and loop forever
for(;;)
  {
	$count++;

	# query the stat table for our data. 
	$q = $db->prepare("select relname, heap_blks_read, heap_blks_hit, idx_blks_read, idx_blks_hit, toast_blks_read, toast_blks_hit, tidx_blks_read, tidx_blks_hit from pg_statio_all_tables");
	
	$q->execute;

	$q->bind_columns(\$tbl, \$heapRead, \$heapHit, \$idxRead, \$idxHit, \$toastRead, \$toastHit, \$toastIdxRead, \$toastIdxHit);

	while($q->fetch)
	  {
		update($tbl, "total", $heapRead + $idxRead + $toastRead + $toastIdxRead, $heapHit + $idxHit + $toastHit + $toastIdxHit);
		update($tbl, "heap", $heapRead, $heapHit);
		update($tbl, "index", $indexRead, $indexHit);
		update($tbl, "toast", $toastRead, $toastHit);
		update($tbl, "toastidx", $toastIdxRead, $toastIdxHit);
	  }

	
	# display results
	@dispOrder = sort { $results{$b}{$sortBy} <=> $results{$a}{$sortBy} } keys %results;
	
	$m = scalar(@dispOrder);
	if($m > 20)
	  {
		$m = 20;
	  }

	printf("\033[2J\033[1;1f");	
	$t = localtime(time);
	print "PG IOMonitor\t$dbname\t".($ioOnly? "IO Only" : "IO+Cache")."\t\t$t\n";
	printf("%10s %14s\t(T)otal\tH(e)ap\t(I)ndex\tT(o)ast\tTo(a)stIdx\n", $unitDisp, "Table");
	printf("--------------------------------------------------------------------------------\n");

	if($count > 1) # don't display the first pass.
	  {
		for($x = 0; $x < $m; $x++)
		  {
			my $tbl = $dispOrder[$x];
			printf("%25s\t%d\t%d\t%d\t%d\t%d\n",  $tbl, $results{$tbl}{total}, 
				   $results{$tbl}{heap}, $results{$tbl}{index}, 
				   $results{$tbl}{toast}, $results{$tbl}{toastidx});
		  }
	  }

	$invec = '';
	vec($invec, fileno(STDIN), 1) = 1;

	($nf, $left) = select($invec, undef, undef, $interval);
	if($nf > 0)
	  {
		$in = getc;
		$in = lc($in);
		# now decide what to do
		if($in eq "q")
		  {
			cleanupAndExit();
			exit;
		  }
		elsif($in eq "e")
		  {
			$sortBy = "heap";
			print "Sort by heap";
		  }
		elsif($in eq "c")
		  {
			$ioOnly = !$ioOnly;
			# reset ourselves so we won't display again.
			$count = 0;
			print "Toggle IO Only";
		  }
		elsif($in eq "d")
		  {
			$sortBy = "index";
			print "Sort by index";
		  }
		elsif($in eq "t")
		  {
			$sortBy = "total";
			print "Sort by total";
		  }
		elsif($in eq "o")
		  {
			$sortBy = "toast";
			print "Sort by toast";
		  }
		elsif($in eq "a")
		  {
			$sortBy = "toastidx";
			print "Sort by toastidx";
		  }
		elsif($in eq "b")
		  {
			if($dispInBlocks)
			  {
				$dispInBlocks = undef;
				$unitDisp = "kB/sec";
				$blockSizeModifier = $blockSize;
				print "Display kB/second";
			  }
			else
			  {
				$dispInBlocks = 1;
				$blockSizeModifier = 1;
				$unitDisp = "blk/sec";
				print "Display blocks / second";
			  }
		  }
		elsif($in eq "h" || $in eq "?")
		  {
			helpScreen();
			next; # skip wait
		  }
		# TODO: select may return nothing in the $left field. or junk
		# it is platform dependant.
		select(undef, undef, undef, $left);
	  }
  }


# update
#
# update our values
#
sub update
  {
	my($tbl, $field, $read, $hit) = @_;
	my($disp, $val);

	$val = $read;
	if(!$ioOnly)
	  {
		$val += $hit;
	  }

	# normalize for interval & block size
	$val = ($val / $interval);

	$disp = ($val - $lastResults{$tbl}{$field}) * $blockSizeModifier;
	$results{$tbl}{$field} = $disp;
	$lastResults{$tbl}{$field} = $val;
  }	

# helpScreen
#
# display the help screen
#
sub helpScreen
  {
	$invec = '';
	vec($invec, fileno(STDIN), 1) = 1;
	
	printf("\033[2J\033[1;1f");	
	print "PG IO Monitor v$version by Jeff Trout <jeff\@jefftrout.com>\n";
	print "http://postgresql.jefftrout.com/\n";
	print "\n";
	print "q - quit\n";
	print "c - Toggle IO / IO + Cache mode\n";
	print "b - Toggle kB/sec vs block/sec units\n";
	print "\n";
	print "Sort Options:\n";
	print "t - sort by Total\n";
	print "e - sort by hEap\n";
	print "i - sort by Index\n";
	print "o - sort by tOast\n";
	print "a - sort by toAst index\n";
	print "\n";
	print "Press any key to return";

	select($invec, undef, undef, undef);
  }

# usage
# 
# display usage
#
sub usage
  {
	print "pgiomonitor [-v version] [-u user - user to connect as] [-p pass - password] [-b display in blocks] [-c show IO only (no cache)] [-t sort by total] [-i sort by index] [-o sort by toast] [-a sort by toast index] db connect string\n";
	print "You can also use \$DBI_DSN & company to set connect options\n";
  }


# setupTerm
#
# turn off canon & echo
sub setupTerm
  {
	# set term attributes so we don't need to hit enter
	$term = new POSIX::Termios;
	$attr = $term->getattr(fileno(STDIN)) or die("getattr failed: $!");
	$lflags = $term->getlflag;
	# turn off canon (don't need to hit enter)
	$lflags &= ~&POSIX::ICANON;
	# turn off echo 
	$lflags &= ~&POSIX::ECHO;
	
	$term->setlflag($lflags);
	$term->setattr(fileno(STDIN), TCSANOW);
  }

# cleanupTerm
# 
# turn on canon & echo
#
sub cleanupAndExit
  {
	$db->disconnect;

	# set term attributes so we don't need to hit enter
	$term = new POSIX::Termios;
	$attr = $term->getattr(fileno(STDIN)) or die("getattr failed: $!");
	$lflags = $term->getlflag;
	# turn off canon (don't need to hit enter)
	$lflags |= &POSIX::ICANON;
	# turn off echo 
	$lflags |= &POSIX::ECHO;
	
	$term->setlflag($lflags);
	$term->setattr(fileno(STDIN), TCSANOW);
 	
	exit;
  }
