#!/usr/bin/perl -w

######################################################################################
# CheckQueryRunTime.pl
# Version 1.02
# Written by: Yuval Baror (http://yuval.bar-or.org)
# Copyright Yuval Baror 2011
######################################################################################

######################################################################################
#    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 3 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.
#
#    See <http://www.gnu.org/licenses/> for a full copy of the
#    GNU General Public License
######################################################################################

### Packages #########################################################################
use strict;
use DBI;
use Time::HiRes qw/gettimeofday tv_interval/;

### Globals ##########################################################################
use vars qw/$dbname $dbuser $dbpass $dbh/;

### Function Declarations ############################################################

### Signal Handling ##################################################################

### Main Function ####################################################################
sub main(@)
{
	my $printResults = 0;
	my $numCacheTimes = 50;
	my $numNoCacheTimes = 10;
	my $runwayTimes = 5;
	
	# Read the commandline arguments
	while (@_ and $_[0] =~ /^-/)
	{
		my $arg = shift;
		if ($arg =~ /^-h(elp)?$/)
		{
			Usage();
			next;
		}
		elsif ($arg =~ /^-p(rintResults)?$/)
		{
			$printResults = 1;
			next;
		}
		elsif ($arg =~ /^-n(umTimes)?$/)
		{
			$numCacheTimes = shift;
			$numNoCacheTimes = $numCacheTimes;
			next;
		}
		elsif ($arg =~ /^-n(c|umCacheTimes)?$/)
		{
			$numCacheTimes = shift;
			next;
		}
		elsif ($arg =~ /^-n(nc|umNoCacheTimes)?$/)
		{
			$numNoCacheTimes = shift;
			next;
		}
		elsif ($arg =~ /^-r(unwayTimes)?$/)
		{
			$runwayTimes = shift;
			next;
		}
		
		warn "Unknown argument '$arg'\n";
		Usage();
	}	
	
	if (scalar(@_) < 4)
	{
		Usage();
	}
	
	my $query;
	($dbname, $dbuser, $dbpass, $query) = @_;
	
	Connect();

	# Run the query with the cache - start by running it 5 times without storing the 
	# result to make sure the data is in the cache
	print "Preparing to run query with cache\n";
	for (my $i = 0; $i < $runwayTimes; ++$i)
	{
		GetQueryTime($query);
	}
	
	print "Running query with cache\n";
	my @cacheTimes = ();
	my $cacheSum = 0;
	my $cacheSumSquares = 0;
	for (my $i = 0; $i < $numCacheTimes; ++$i)
	{
		print ".";
		my $result = GetQueryTime($query);
		if ($result)
		{
			push(@cacheTimes, $result);
			$cacheSum += $result;
			$cacheSumSquares += $result * $result;
		}
	}
	print "\n";
	
	my $cacheAverage = "Invalid";
	my $cacheStddev = "Invalid";
	my $cacheCount = scalar(@cacheTimes);
	if ($cacheCount > 0)
	{
		$cacheAverage = $cacheSum / $cacheCount;
		my $cacheVariance = $cacheSumSquares/$cacheCount - $cacheAverage*$cacheAverage;
		$cacheStddev = ($cacheVariance < 0 ? 0 : sqrt($cacheVariance));
	}
	
	print "With cache: $cacheAverage ms. (stddev $cacheStddev ms.)\n";
	if ($printResults)
	{
		print "[" . join(", ", @cacheTimes) . "]\n";
	}
	
	# Run the query without the cache
	print "\nRunning query without cache\n";
	
	my @noCacheTimes = ();
	my $noCacheSum = 0;
	my $noCacheSumSquares = 0;
	for (my $i = 0; $i < $numNoCacheTimes; ++$i)
	{
		print ".";
		ClearCache();
		my $result = GetQueryTime($query);
		if ($result)
		{
			push(@noCacheTimes, $result);
			$noCacheSum += $result;
			$noCacheSumSquares += $result * $result;
		}
	}
	print "\n";
	
	my $noCacheAverage = "Invalid";
	my $noCacheStddev = "Invalid";
	my $noCacheCount = scalar(@noCacheTimes);
	if ($noCacheCount > 0)
	{
		$noCacheAverage = $noCacheSum / $noCacheCount;
		my $noCacheVariance = $noCacheSumSquares/$noCacheCount - $noCacheAverage*$noCacheAverage;
		$noCacheStddev = ($noCacheVariance < 0 ? 0 : sqrt($noCacheVariance));
	}
	
	print "Without cache: $noCacheAverage ms. (stddev $noCacheStddev ms.)\n";
	if ($printResults)
	{
		print "[" . join(", ", @noCacheTimes) . "]\n";
	}
	
	Disconnect();
	
	return 0;
}

exit(main(@ARGV));

######################################################################################
### Subroutines section 
######################################################################################

######################################################################################
# Name: Usage
# Purpose: Print a usage message and quit the program
# Arguments: none
# Return Value: none
# Semantics: none
######################################################################################
sub Usage()
{
	die "Usage: $0 [-printResults] [-numTimes <num>] [-numCacheTimes <num>] [-numNoCacheTimes <num>] [-runwayTimes <num>] <dbName> <dbUsername> <dbPassword> <query>\n\n" . 
	"\t-printResults: print the run times of all queries instead of just summaries\n" .
	"\t-numTimes: the number of times to run the query both with the cache and without the cache\n" .
	"\t-numCacheTimes: the number of times to run the query with the cache\n" .
	"\t-numNoCacheTimes: the number of times to run the query without the cache\n" .
	"\t-runwayTimes: the number of times to run the query to make sure its data is in the cache\n" .
	"\tdbName: the name of the database to connect to\n" .
	"\tdbUsername: the username to access the database with\n" .
	"\tdbPassword: the password to access the database with\n" .
	"\tquery: the query to test\n" .
	"\n\nFor example: $0 mydb myuser mypass \"select * from my_table where my_column=17\"\n";
}

# Connect to the DB
sub Connect()
{
	$dbh = DBI->connect("DBI:Pg:dbname=$dbname", $dbuser, $dbpass, {AutoCommit => 1, RaiseError => 0});
	if (!$dbh) {
		die "Cannot connect to Postgres server: $DBI::errstr\n";
	}
}

# Disconnect from the DB
sub Disconnect()
{
	$dbh->disconnect();
	undef $dbh;
}

# Clear the DB cache (actually clearing the OS caches)
sub ClearCache()
{
	# Disconnect from the DB
	Disconnect();
	
	# Clear the caches
	system("service postgresql stop &> /dev/null; sync; echo 3 > /proc/sys/vm/drop_caches; service postgresql start &> /dev/null");
	
	# Reconnect to the DB
	Connect();
}

# Get the run time of a query in milliseconds
sub GetQueryTime($)
{
	my $query = shift;
	my @startTime = gettimeofday();
	my $result = $dbh->do($query);
	my @endTime = gettimeofday();
	if (!$result)
	{
		print "Error running query: $query\n";
		print $dbh->errstr() . "\n";
		return 0;
	}
	
	return tv_interval(\@startTime, \@endTime) * 1000;
}