#!/bin/bash

#
# pgsqlslower NMILLIS: report details on slow postgres queries
#

ps_arg0="pgsqlslower"

#
# Some of this boilerplate is common to several tools, but it's duplicated here
# because it's very useful for these to be standalone scripts.
#
ps_tmpfile="/var/tmp/$ps_arg0.$$"
ps_number_re='^[1-9][0-9]*$'
ps_synopsis="usage: $ps_arg0 NMILLISECONDS"

if [[ $# == 0 ]]; then
	echo "$ps_synopsis" >&2
	exit 2
fi

if [[ "$1" == "-h" || "$1" == "--help" ]]; then
	cat >&2 <<EOF
$ps_synopsis

Print details about queries taking longer than NMILLISECONDS from start to
finish on all postgresql instances on this system.  Note that since this tool
traces query start to query done, it will never see queries taking longer than
NMILLISECONDS seconds.

This is similar to pgsqltxslower, except that it's tracing simple queries.  For
extended queries that are part of transactions, see pgsqltxslower.

This tool requires privileges to use DTrace on postgres processes on this
system.  If you see no output but expect some, check whether your user has
permissions to trace the postgres processes.
EOF
	exit 2
elif ! [[ "$1" == "0" || "$1" =~ $ps_number_re ]]; then
	echo "$ps_arg0: bad number of milliseconds" >&2
	echo "$ps_synopsis" >&2
	exit 2
fi

trap cleanup EXIT
function cleanup
{
	rm -f "$ps_tmpfile"
}

if ! type dtrace > /dev/null 2>&1; then
	echo "$ps_arg0: requires dtrace(1M), but not found" >&2
	exit 1
fi

cat > "$ps_tmpfile" <<EOF
#!/usr/sbin/dtrace -Cs

#pragma D option quiet
#pragma D option zdefs
#pragma D option strsize=1024

#define	THRESHOLD ($1)
#define	MICROS(X)	((X) / 1000)
#define	MICRODIFF(X, Y)	(MICROS((X) - (Y)))

/*
 * XXX can queries be started recursively?  If so, this will do the right thing
 * on query entry (only capturing the top-level one), but not on query-done.
 */
postgresql*:::query-start
/!self->start/
{
	self->start = timestamp;
	self->qstr = arg0;

	self->nreads = 0;
	self->ncachehits = 0;
	self->nflushes = 0;

	self->parse_start = 0;
	self->parse_done = 0;
	self->plan_start = 0;
	self->plan_done = 0;
	self->exec_start = 0;
	self->exec_done = 0;

	self->ntxnstarts = 0;
	self->ntxncommits = 0;
	self->ntxnaborts = 0;
}

/*
 * Track elapsed time during query processing phases.
 */
postgresql*:::query-parse-start
/self->start/
{
	self->parse_start = timestamp;
}

postgresql*:::query-parse-done
/self->parse_start/
{
	self->parse_done = timestamp;
}

postgresql*:::query-plan-start
/self->start/
{
	self->plan_start = timestamp;
}

postgresql*:::query-plan-done
/self->plan_start/
{
	self->plan_done = timestamp;
}

postgresql*:::query-execute-start
/self->start/
{
	self->exec_start = timestamp;
}

postgresql*:::query-execute-done
/self->exec_start/
{
	self->exec_done = timestamp;
}

/*
 * Track buffer reads.
 */
postgresql*:::buffer-read-done
/self->start && arg7/
{
	self->ncachehits++;
}

postgresql*:::buffer-read-done
/self->start/
{
	self->nreads++;
}

postgresql*:::buffer-flush-done
/self->start/
{
	self->nflushes++;
}

/*
 * Track transactions started, committed, or aborted.
 */
postgresql*:::transaction-start
/self->start/
{
	self->ntxnstarts++;
}

postgresql*:::transaction-commit
/self->start/
{
	self->ntxncommits++;
}

/*
 * On a transaction abort, the current query is done, and we won't see a
 * query-done probe.
 */
postgresql*:::transaction-abort
/self->start/
{
	self->ntxnaborts++;
}

postgresql*:::transaction-abort,
postgresql*:::query-done
/self->start && (timestamp - self->start) > THRESHOLD * 1000000/
{
	printf("QUERY: %s\n", copyinstr(self->qstr));
	printf("   total time: %8d us (parse/plan/execute = %dus/%dus/%dus)\n",
	    MICRODIFF(timestamp, self->start),
	    self->parse_done == 0 ? 0 :
	        MICRODIFF(self->parse_done, self->parse_start),
	    self->plan_done == 0 ? 0 :
	        MICRODIFF(self->plan_done, self->plan_start),
	    self->exec_done == 0 ? 0 :
	        MICRODIFF(self->exec_done, self->exec_start));
	printf("         txns: %d started, %d committed, %d aborted\n",
	    self->ntxnstarts, self->ntxncommits, self->ntxnaborts);
	printf("      buffers: %d read (%d hit, %d missed), %d flushed\n",
	    self->nreads, self->ncachehits, self->nreads - self->ncachehits,
	    self->nflushes);
	printf("\n");
}

/*
 * Clean up state.
 */
postgresql*:::transaction-abort,
postgresql*:::query-done
{
	self->qstr = 0;
	self->start = 0;
	self->nreads = 0;
	self->ncachehits = 0;
	self->nflushes = 0;
	self->parse_start = 0;
	self->parse_done = 0;
	self->plan_start = 0;
	self->plan_done = 0;
	self->exec_start = 0;
	self->exec_done = 0;
	self->ntxnstarts = 0;
	self->ntxncommits = 0;
	self->ntxnaborts = 0;
}
EOF

dtrace -Cs "$ps_tmpfile"