/* RT monitor.

        Copyright (c) 2002 Roger Larsson <roger.larsson_AT_norran.net>

    This program is free software; you can redistribute it and/or
    modify it under the terms of version 2 of the GNU General Public
    License as published by the Free Software Foundation.

    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.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.

    Thanks to autor of KSysGuard Chris Schlaeger for borrowed code...
*/

#include <sys/types.h>
#include <sched.h>
#include <stdio.h>
#include <dirent.h>
#include <sys/mman.h>

 int isRT(pid_t pid)
 {
     int sched_class = sched_getscheduler( pid);
     if (sched_class == -1) {
	 fprintf(stderr, "Pid %d Exited?\n", pid);
	 return 0;
     }

     return sched_class != SCHED_OTHER;
 }

struct rt_process_info
{
    /* This flag is set for all found processes at the beginning of the
     * process list update. Processes that do not have this flag set will
     * be assumed dead and removed from the list. The flag is cleared after
     * each list update. */
    int alive;
    int centStamp;


    pid_t pid;
    pid_t ppid;
    gid_t gid;

    unsigned int userTime;
    unsigned int sysTime;
    unsigned int vmSize; // enough?
    unsigned int vmRss; // enough?

    float sysLoad;
    float userLoad;
    float cpu_usage;
};

#define MAX_RT_PROCESSES 200
struct rt_process_info rt_process[MAX_RT_PROCESSES]; /* pid & alive == 0 */

struct rt_process_info *find_process(pid_t pid)
{
    unsigned ix;
    for (ix = 0; ix < MAX_RT_PROCESSES; ix++)
    {
	if (rt_process[ix].pid == pid) {
	    rt_process[ix].alive = 1;
	    return &rt_process[ix];
	}
    }

    return NULL;
}

struct rt_process_info *new_process(pid_t pid)
{
    unsigned ix;
    for (ix = 0; ix < MAX_RT_PROCESSES; ix++)
    {
	if (rt_process[ix].pid == 0) {
	    rt_process[ix].pid = pid;
	    rt_process[ix].alive = 2;
	    return &rt_process[ix];
	}
    }

    return NULL;
}

float cpu_usage(struct rt_process_info *ps)
{
#define BUFSIZE 1024
    char buf[BUFSIZE];
    FILE *fd;
    char status;
    unsigned int userTime, sysTime;

    snprintf(buf, BUFSIZE - 1, "/proc/%d/stat", ps->pid);
    buf[BUFSIZE - 1] = '\0';
    if ((fd = fopen(buf, "r")) == 0)
	return (-1);

    if (fscanf(fd, "%*d %*s %c %d %d %*d %*d %*d %*u %*u %*u %*u %*u %d %d"
	       "%*d %*d %*d %*d %*u %*u %*d %u %u",
	       &status, (int*) &ps->ppid, (int*) &ps->gid,
	       &userTime, &sysTime, &ps->vmSize,
	       &ps->vmRss) != 7) {
	fclose(fd);
	return (-1);
    }

    if (fclose(fd))
	return (-1);

    {
	unsigned int newCentStamp;
	int timeDiff, userDiff, sysDiff;
	struct timeval tv;

	gettimeofday(&tv, 0);
	newCentStamp = tv.tv_sec * 100 + tv.tv_usec / 10000;

	// calculate load
	if (ps->alive == 2)
	    ps->sysLoad = ps->userLoad = 0.0f; /* can't give relieable number at the moment... */
	else {
	    timeDiff = (int)(newCentStamp - ps->centStamp);
	    userDiff = userTime - ps->userTime;
	    sysDiff = sysTime - ps->sysTime;

			
	    if ((timeDiff > 0) && (userDiff >= 0) && (sysDiff >= 0)) /* protect from bad data */
	    {
		ps->userLoad = ((double) userDiff / timeDiff) * 100.0;
		ps->sysLoad = ((double) sysDiff / timeDiff) * 100.0;
	    }
	    else
		ps->sysLoad = ps->userLoad = 0.0;
	}

	// update fields
	ps->centStamp = newCentStamp;
	ps->userTime = userTime;
	ps->sysTime = sysTime;
    }
	
    ps->cpu_usage = ps->userLoad + ps->sysLoad;

    return ps->cpu_usage;
}

/* process reading code from ksysguard */
float cpu_rt_usage(struct rt_process_info **rt_list_head)
{
    // Watch out for SMP effects...
	
    float result = 0.0f;
    pid_t myself = getpid();
    DIR* dir;
    struct dirent* entry;

    /* read in current process list via the /proc filesystem entry */
    if ((dir = opendir("/proc")) == NULL)
    {
	perror("Cannot open directory \'/proc\'!\n"
	       "The kernel needs to be compiled with support\n"
	       "for /proc filesystem enabled!\n");
	return 0;
    }
	
    // for all processes
    while ((entry = readdir(dir)))
    {
	if (isdigit(entry->d_name[0]))
	{
	    pid_t pid;

	    pid = atoi(entry->d_name);
			
	    if (pid != myself && isRT(pid)) {
		struct rt_process_info *process = find_process(pid);
		float cpu_use;

		printf("Found a RT process %d, info 0x%x\n", pid, process);

		if (process == NULL)
		{
		    process = new_process(pid);
		    if (process == NULL) {
			// to many RT processes!
			//  this process is new - assume a DOS attack
			printf("Out of RT process info space - "
			       "assume DOS attack\n");
			set_normal_priority(pid);
		    }

						
		    process->alive = 2; /* mark process new */
		}

		cpu_use = cpu_usage(process);

		process->alive = 1;
		result += cpu_use;
	    }
	}
    }
    closedir(dir);
	
    return result;
}


void gc_rt_processes()
{
    unsigned ix;
    for (ix = 0; ix < MAX_RT_PROCESSES; ix++)
    {
	struct rt_process_info *rt_examine = &rt_process[ix];

	if (rt_examine->alive)
	{
	    rt_examine->alive = 0;
	}
	else
	{
	    rt_examine->pid = 0; /* delete it! */
	}
    }
}

int set_realtime_priority(void)
{
struct sched_param schp;
	/*
	 * set the process to realtime privs
	 */
        memset(&schp, 0, sizeof(schp));
	schp.sched_priority = sched_get_priority_max(SCHED_FIFO);

	if (sched_setscheduler(0, SCHED_FIFO, &schp) != 0) {
		perror("sched_setscheduler");
		return -1;
	}

	if(mlockall(MCL_CURRENT|MCL_FUTURE))
	{
	    perror("mlockall() failed, exiting. mlock");
	    return -1;
	}

	return 0;

}

int set_normal_priority(pid_t pid)
{
struct sched_param schp;
	/*
	 * set the process to realtime privs
	 */
        memset(&schp, 0, sizeof(schp));
	schp.sched_priority = 0;

	printf("Attempt to reduce scheduling class for pid %d ", pid);
	if (sched_setscheduler(pid, SCHED_OTHER, &schp) != 0) {
		printf("- failed!\n");
		perror("sched_setscheduler");
		return -1;
	}
	printf("- done!\n");

	return 0;
}

void set_normal_priority_all()
{
    unsigned ix;
    for (ix = 0; ix < MAX_RT_PROCESSES; ix++)
    {
	struct rt_process_info *process_info = &rt_process[ix];
	if (process_info->pid)
	    set_normal_priority(process_info->pid);
    }
}

int cpu_idle()
{
    static int _user, _sys, _nice, _idle;
    int user, sys, nice, idle;
    int possible_rt_work, non_rt_work;

    // read file /proc/stat (first line: user, sys, nice, idle)
    FILE *stat = fopen("/proc/stat", "r");
    fscanf(stat, "%*s %d %d %d %d", &user, &sys, &nice, &idle);
    fclose(stat);

    possible_rt_work = (user - _user + sys - _sys);
    non_rt_work = (nice - _nice + idle - _idle);
    _user = user; _sys = sys; _nice = nice; _idle = idle;

    return 100 * non_rt_work / (non_rt_work + possible_rt_work);
}

int main(int argc, char * argv[])
{
    struct rt_process_info *rt_list = NULL;
    
    // monitor process runs with realtime prio
    set_realtime_priority();
 #define MIN_IDLE 10
 #define MAX_RT_USAGE 70

    while (1) {
	if (cpu_idle() < MIN_IDLE) {
	    printf("Total CPU IDLE below MIN_IDLE\n");
	    if (cpu_rt_usage(&rt_list) > MAX_RT_USAGE) {
		printf("Total CPU RT usage above MAX_RT_USAGE\n");

		gc_rt_processes();

		// build process trees from rt_list
		// decide which tree to reduce to normal prio class
		//   (assume only one for simplicitly...)

		// reduce all processes in that tree
		set_normal_priority_all();

		//   (may use nice to simulate prio levels)
		// log a message

	    }		
	}
	sleep(10);
    }
    
    // process exiting - free elements on rt_list...
    //free_rt_list(rt_list);
}