/*
 *  resched (Roman Drahtmueller <draht@uni-freiburg.de>
 *
 *  delivers a process to another scheduling class policy.
 *
 *  CAUTION: Handle with care and read the manpage before you use it.
 *           CPU-intesive processes will not return before they are
 *           finished or blocked on I/O. (This means that it could 
 *           hang your machine, luser!)
 *
 *  Use as you wish. The author is not responsible for
 *  damages that could result from using this software.
 *  Permission granted to distribute at will.
 *
 */


/*	#define SCHED_OTHER     0
	#define SCHED_FIFO      1
	#define SCHED_RR        2
	#define SCHED_IDLE	5
*/

#define MAX( a, b ) (a>b ? a : b)
#define MIN( a, b ) (a>b ? b : a)
#define PROGNAME "resched"
/* #define RENICE_ONLY_FOR_SCHED_OTHER 1 */

#define TOPLINE \
  "  pid   OLD policy   priority nice |   NEW policy    priority        nice\n" \
  "-------------------------------------------------------------------------\n"

#define TOPLINE_LISTONLY \
  "  pid       policy   priority nice\n" \
  "----------------------------------\n"

#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <sched.h>
#include <linux/sched.h>

char policy_name[][12] = { "SCHED_OTHER", "SCHED_FIFO", "SCHED_RR", "SCHED_BATCH", "NONE", "SCHED_IDLE", "" };

extern int atoi();
void usage(void) {
            fprintf(stderr,
"%s is a utility to change and see the scheduling policy, the priority\n"
"and the nicelevel for one or more processes.\n"
"Usage: a) %s [-p <priority>] [-o|-f|-r|-i|-b] [-n <nice-value>] <pid> ...\n"
"       b) %s -l <pid> ...               # this is the default\n"
"       c) %s <>\n"
"       d) %s -h\n\n", PROGNAME, PROGNAME, PROGNAME, PROGNAME, PROGNAME);

	fprintf(stderr,
"a) %7s will change the scheduling policy of the process(es) <pid> ...\n"
"        -o means to set SCHED_OTHER\n"
"        -f is SCHED_FIFO\n"
"        -r is SCHED_RR\n"
"        -b is SCHED_BATCH\n"
"        -i is SCHED_IDLE\n"
"        -p <priority> adds the realtime schedulung priority, which must be\n"
"           zero for SCHED_OTHER.\n"
"        -n <nice-value> is equivalent to the -n option of the nice(1)\n"
"           command, changing the nice value for SCHED_OTHER scheduling policy.\n"
"b) list mode. %s will not change any processes. This is the default.\n"
"c) without any arguments, %s will tell you about its current scheduling\n"
"   class, the corresponding priority and the nice value in SCHED_OTHER\n"
"d) the -h commandline option shows this brief help text.\n",
	PROGNAME, PROGNAME, PROGNAME);
	fprintf(stderr,
"%s returns 0 upon sucess or the number of failures during its operation\n",
	PROGNAME);
	fprintf(stderr, "%s",
"To learn more about scheduling policies, please see the manpages for\n"
"   sched_setscheduler(2), sched_setparam(2), getpriority(2), nice(1)\n"
"and the SEE ALSO hints in these manpages.\n\n"
"---===   WARNING:   ===---\n"
"Delivering processes (one per processor or more) to SCHED_RR or SCHED_FIFO\n"
"realtime scheduling policy may cause your system to freeze (without being\n"
"crashed!)!\n"
"If you are new to this utility and if you want to play with it, then\n"
"run the command (from the same shell as you want to test with)\n"
"\tulimit -t 20\nto limit the CPU time of your testbed to 20 seconds.\n");
	fprintf(stderr,
"\n%s is released under the license terms of the Gnu General Public License.\n"
"Use as you wish, but on your own risk.\n"
"Please report bugs to Roman Drahtmueller <draht@suse.de>.\n", PROGNAME);
}

int main(int argc, char **argv) {

char c;
pid_t pid;
struct sched_param *priority;
struct sched_param *old_priority;
int set_priority=0;

int policy = -1;
int old_policy;
int set_policy=0;

int nice;
int old_nice;
int set_nice=0;

int listonly=0;
int err=0;


pid=0;

    priority = (struct sched_param *) malloc(sizeof(struct sched_param));
    priority->sched_priority = 0;
    old_priority = (struct sched_param *) malloc(sizeof(struct sched_param));


    while ((c = getopt(argc, argv, "p:forhln:ib")) != EOF)
        switch (c) {
        case 'p':
            priority->sched_priority = atoi(optarg);
	    set_priority=1;
            break;
        case 'o':
            policy = SCHED_OTHER;
	    if(set_policy==1) {
		usage();
		exit(-1);
	    }
	    set_policy=1;
            break;
        case 'f':
            policy = SCHED_FIFO;
            if(set_policy==1) {
                usage();
                exit(-1);
            }
	    set_policy=1;
            break;
        case 'r':
            policy = SCHED_RR;
            if(set_policy==1) {
                usage();
                exit(-1);
            }
	    set_policy=1;
            break;
        case 'b':
            policy = SCHED_BATCH;
            if(set_policy==1) {
                usage();
                exit(-1);
            }
            set_policy=1;
            break;
        case 'i':
            policy = SCHED_IDLE;
            if(set_policy==1) {
                usage();
                exit(-1);
            }
            set_policy=1;
            break;
	case 'l':
	    if(set_policy==1) {
		usage();
		exit(-1);
	    }
	    listonly = 1;
	    break;
        case 'n':
	    nice=atoi(optarg);
	    set_nice=1;
	    break;
        case 'h':
            usage();
            exit (-1);
        default:
	    usage();
	    exit(-1);
            break;
        }


    if( (set_priority==0) && (set_policy==0) && (set_nice==0) && (set_nice==0))
	listonly=1;

    if(argc-optind==0) { /* don't have a pid on the commandline. list 
			    self only and spit out a comment. */
	policy=sched_getscheduler(pid);
	sched_getparam(pid, old_priority);
	old_nice=getpriority(PRIO_PROCESS, pid);

	printf("...running under scheduling policy %i (%s), prio %i "
		"and nice level %i.\n"
		"Type \"%s -h\" for help.\n",
		policy, policy_name[policy], old_priority->sched_priority,
			old_nice, PROGNAME);
	exit(0);
    }

    printf("%s", listonly ? TOPLINE_LISTONLY : TOPLINE);


    /* loop over PIDs on the commandline. */
    while( (argc - optind) >=1) {

	pid=atoi(argv[optind]);


            if(sched_getparam(pid, old_priority)) {
                fprintf(stderr, "%s: %i: %s\n", PROGNAME, pid, strerror(errno));
                err++;
                goto err_and_next;
            }
	    if( (old_policy=sched_getscheduler(pid)) <0) {
                fprintf(stderr, "%s: %i: %s\n", PROGNAME, pid, strerror(errno));
                err++;
                goto err_and_next;
            }

	    errno=0;
	    old_nice=getpriority(PRIO_PROCESS, pid);
	    if( (old_nice==-1) && (errno!=0)) {
		fprintf(stderr,"%s: getpriority(): %s", PROGNAME, strerror(errno));
		err++;
                goto err_and_next;
            }


            printf("%5.0i  %11s     %6i  %3i ", 
		pid, policy_name[old_policy], old_priority->sched_priority, old_nice);
	if(listonly)
	    goto set_nice;
	else {

/* ------------ done listing. */

	    if(set_priority && set_policy) {
    		if(policy == SCHED_OTHER)
		    if(priority->sched_priority) { /* wrong -p value on the commandline. */
			priority->sched_priority=0;
			fprintf(stderr, "%s: ignoring bad priority for "
					"SCHED_OTHER scheduling policy.\n",PROGNAME);
		    }
		/* allowed values: */
        	priority->sched_priority=MIN(priority->sched_priority, 
					sched_get_priority_max(policy));
	        priority->sched_priority=MAX(priority->sched_priority, 
					sched_get_priority_min(policy));

                if (sched_setscheduler(pid,policy,priority)) {
                    fprintf(stderr,"%s: sched_setscheduler(): %s\n", PROGNAME, strerror(errno));
                    err++;
                    goto err_and_next;
		}
		if(old_policy==policy) {
		    printf("|  %11s","unchanged");
		} else {
		    printf("|  %11s", policy_name[policy]);
		}
		if (old_priority->sched_priority == priority->sched_priority) {
		    printf(" %11s","unchanged");
		} else {
		    printf(" %11i",priority->sched_priority);
		}
	    }


	    if(set_priority && !set_policy) {
		if( (old_policy==SCHED_OTHER) && 
			(priority->sched_priority!=0)) {
			fprintf(stderr, "%s: ignoring priority change for "
					"SCHED_OTHER policy process\n",PROGNAME);
			/* don't count this as an error. And goto is ugly, yes. */
/*			goto set_nice; */
		}
        	priority->sched_priority=MIN(priority->sched_priority, 
					sched_get_priority_max(old_policy));
	        priority->sched_priority=MAX(priority->sched_priority, 
					sched_get_priority_min(old_policy));
		if(old_priority->sched_priority == priority->sched_priority) {
			printf("|  %11s %11s", "unchanged","unchanged");
		} else if(sched_setparam(pid,priority)) {
			fprintf(stderr,"%s: sched_getparam(): %s\n",PROGNAME, strerror(errno));
			err++;
			goto err_and_next;
		} else {
			printf("|  %11s %6i", "unchanged",
							priority->sched_priority);
		}
	    }

	    if(!set_priority && set_policy) {
		if(old_policy == policy) {
		    printf("|  %11s %11s","unchanged", "unchanged");
		    goto set_nice;
		}
		if(policy != SCHED_OTHER) {
		    priority->sched_priority=old_priority->sched_priority;
		    priority->sched_priority=MIN(priority->sched_priority,
					sched_get_priority_max(policy));
		    priority->sched_priority=MAX(priority->sched_priority,
					sched_get_priority_min(policy));
		} else {
		    priority->sched_priority=0;
		}

		if(sched_setscheduler(pid,policy,priority)) {
		    fprintf(stderr,"%s: sched_setscheduler(): %s\n", PROGNAME, strerror(errno));
		    err++;
		    goto err_and_next;
		}

		if(priority->sched_priority != old_priority->sched_priority) {
		    printf("|  %11s %11i",
			policy_name[policy], priority->sched_priority);
		} else {
		    printf("|  %11s %11s", policy_name[policy], "unchanged");
		}
	    }

	    if(!set_priority && !set_policy) {
		printf("|  %11s %11s","unchanged", "unchanged");
	    }

			/* change nice levels only for SCHED_OTHER processes: */
set_nice:
#ifdef RENICE_ONLY_FOR_SCHED_OTHER
	    if( (set_policy && (policy==SCHED_OTHER)) ||
			(!set_policy && (old_policy==SCHED_OTHER)) ) {
#endif
	        if(! set_nice || (old_nice==nice)) {
		    if(!listonly) 
			printf(" %11s","unchanged");
		} else {
		    if(setpriority(PRIO_PROCESS,pid,nice) < 0) {
		        fprintf(stderr, "%s: setpriority(): %s\n", PROGNAME,strerror(errno));
			err++;
			goto err_and_next;
		    } else {
			printf(" %11i",nice);
		    }
		}
#ifdef RENICE_ONLY_FOR_SCHED_OTHER
	    } else {
		printf(" %11s","unchanged");
	    }
#endif
	    printf("\n");

	}
    goto do_next;
err_and_next:
    printf("\n");
do_next:
    optind++;
    }


    return(err);

}