#include <sys/types.h>
#include <sys/stat.h>
#include <sys/resource.h>
#include <sys/signal.h>
#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <sched.h>

/** pipe between the two processes **/
int comm[2];
pid_t watch_pid, program_pid;
int result = 1;

int set_realtime(pid_t pid, long preferred)
{
	struct sched_param sp;
	if(preferred != 0)
	{
		long max_pri = sched_get_priority_max(SCHED_FIFO);
		long min_pri = sched_get_priority_min(SCHED_FIFO);
		long priority = (max_pri+min_pri)/2 + preferred - 1;

		if(priority > max_pri) priority = max_pri;

		sp.sched_priority = priority;
	}
	else
	{
		sp.sched_priority = 0;
	}

	if(sched_setscheduler(pid, preferred?SCHED_FIFO:SCHED_OTHER, &sp) != 0)
	{
		perror("Error setting realtime priority: ");
		return 0;
	}
	return 1;
}

void make_nonblocking_pipe()
{
	int i;

	if(pipe(comm) != 0)
	{
		perror("Aborting. Can't create pipe: ");
		exit(1);
	}

	for(i = 0; i < 2; i++)
	{
		int options = fcntl(comm[i], F_GETFL);
		if (options == -1)
		{
			perror("Aborting. Can't make pipe non-blocking: ");
			exit(1);
		}

		if (fcntl(comm[i], F_SETFL, options | O_NONBLOCK) == -1)
		{
			perror("Aborting. Can't make pipe non-blocking: ");
			exit(1);
		}
	}
}

static void exit_cleanup(int sig)
{
	if(program_pid > 0)
	{
		fprintf(stderr,"Removing realtime rights.\n");
		set_realtime(program_pid,0);
	}

	close(comm[0]);
	exit(result);
}

static void init_signals()
{
    signal(SIGHUP ,exit_cleanup);
    signal(SIGQUIT,exit_cleanup);
    signal(SIGINT ,exit_cleanup);
    signal(SIGTERM,exit_cleanup);
}

/**
 * Usage: realtime_set <pid>
 *
 * Install suid root. This gives a program realtime rights, and removes them
 * again, if the system becomes frozen.
 *
 * Problems: It does only globally check if the system freezes, so it will
 * also remove the realtime rights from <pid> if the system is frozen for
 * some other reason (e.g. another realtime program).
 */
int main(int argc, char **argv)
{
	if(argc != 2)
	{
		fprintf(stderr,"usage: %s <pid>\n", argv[0]);
		exit(1);
	}
	make_nonblocking_pipe(comm);

	program_pid = atol(argv[1]);
	watch_pid = fork();
	if(watch_pid == 0)
		/* child process: write dots to comm[0] to indicate we're alive */
	{
		close(comm[0]);
		while(write(comm[1],".",1) == 1)
			sleep(2);
		exit(0);
	}
	else if(watch_pid > 0)
		/* parent process: check dots from comm[1] to see if we have lockups */
	{
		char c;
		int alive = 10;

		if(!set_realtime(getpid(),2))
			exit_cleanup(0);
		if(!set_realtime(program_pid,1))
			exit_cleanup(0);

		for(alive = 30; alive > 0; alive--)
		{
			sleep(1);
			if(read(comm[0],&c,1) == 1)
				alive = 30;

			if(kill(program_pid,0) == -1)
				alive = -1;
		}
		if(alive < 0)
		{
			fprintf(stderr,"Program exited.\n");
			program_pid = 0;
			result = 0;
		}
		exit_cleanup(0);
	}
	else
	{
		perror("Fork failed : ");
	}
}