#!/bin/bash
#===============================================================================
#
#          FILE: pg_supernice
#
#         USAGE: pg_supernice [Directory to locate pg_supernice.so]
#
#   DESCRIPTION: Creating a nice function in C for lowering process priority to 
#                postmaster(server-side process of PostgreSQL) which you are 
#                connected.
#
#                This can makes your query has lowest I/O priority either by 
#                changing the I/O priority class to idle class.
#
#                Idle class means the process will only get disk time when there
#                is no other process asked for disk I/O, so any activity with 
#                the process, the impact should be zero.
#
#  REQUIREMENTS: gcc, pg_config, postgresql-devel
#
#         NOTES: For root user, postmaster has no permission to reach your $HOME.
#                So you will have to specify a public directory, for example:
#                # pg_supernice /usr/local/lib
#
#                pg_supernice require kernel version >= 2.6.25, otherwise the I/O
#                priority change will not work, and process priority will not 
#                effectively.
#
#          BUGS:  ---
#        AUTHOR: rickz (Rick Zhang), xlrickz@gmail.com
#       COMPANY: X-LEGEND Entertainment Corp.
#       CREATED: Wed Aug  6 13:12:19 ICT 2014
#      REVISION: 1.0
#
#          TODO:
#
#===============================================================================

#set -o nounset                              # Treat unset variables as an error
#set -m                                      # Enable job control

# Build directory
[ "$1" == "" ] && BUILD_DIR="$HOME/bin/pg_supernice_so" || BUILD_DIR="$1/pg_supernice"
[ "$USER" == "root" ] && BUILD_DIR="$HOME/bin/pg_supernice_so"
[ "$1" == "" ] && [ "$USER" != "root" ] && chmod 755 $HOME

mkdir -p "$BUILD_DIR"
cd "$BUILD_DIR"

# Source file
cat > pg_supernice.c << __EOF__
#include <sys/time.h>
#include <sys/resource.h>
#include <sys/syscall.h>
#include <unistd.h>
#include "postgres.h"
#include "fmgr.h"

static inline int ioprio_set(int which, int who, int ioprio)
{
        return syscall(SYS_ioprio_set, which, who, ioprio);
}

// For PostgreSQL 8.2+
#ifdef PG_MODULE_MAGIC
PG_MODULE_MAGIC;
#endif

PG_FUNCTION_INFO_V1(pg_supernice);

Datum
pg_supernice(PG_FUNCTION_ARGS)
{
    int result_proc = setpriority(PRIO_PROCESS, 0, 19);
    int result_io = ioprio_set(0x1, 0, 0x6007);

    if (result_proc == 0 && result_io == 0) {
        PG_RETURN_BOOL(1);
    } else {
        PG_RETURN_BOOL(0);
    }
}
__EOF__

# Check compiler and pg_config binary
echo "Checking gcc and pg_config ... "
which gcc || CAN_NOT_COMPILE=1
if which pg_config ; then 
   PG_CONFIG="$(which pg_config)"
else # If pg_config not found
   echo -n "Trying to find pg_config ... "
   PG_CONFIG="$(rpm -ql $(rpm -qa | grep postgresql | grep -v debuginfo)|grep "/bin/pg_config$")"
   echo $PG_CONFIG
   echo -n "Testing pg_config ... "
   $PG_CONFIG | grep VERSION || CAN_NOT_COMPILE=1
fi
echo

# Compile
if [ "$CAN_NOT_COMPILE" == "" ] ; then
   echo Compiling ...
   rm -f pg_supernice.so
   gcc -Wall -I `$PG_CONFIG --includedir-server` -fpic -c pg_supernice.c
   gcc -shared -o pg_supernice.so pg_supernice.o
fi

[ -e pg_supernice.so ] && echo Compiled successfully || echo Compile failed. Use pre-built shared library.

chmod 755 pg_supernice.so

echo

# Create Function pg_supernice(int)
echo "Creating function for databases ... "
for DB in $(psql -U postgres <<< "copy (select datname from pg_database where datname not like 'template%') to stdout csv;") ; do
	echo -en "                            $DB: " | tail -c 28
        psql -U postgres $DB <<< "CREATE OR REPLACE FUNCTION pg_supernice() RETURNS VOID AS '$BUILD_DIR/pg_supernice' LANGUAGE C RETURNS NULL ON NULL INPUT;"
done
