New upstream release.
[debian/cpulimit.git] / cpulimit.c.new
1 /**
2  * This program is licensed under the GNU General Public License,
3  * version 2. A copy of the license can be found in the accompanying
4  * LICENSE file.
5  *
6  **********************************************************************
7  *
8  * Simple program to limit the cpu usage of a process
9  * If you modify this code, send me a copy please
10  *
11  * Author:  Angelo Marletta
12  * Date:    26/06/2005
13  * Version: 1.1
14  *
15  * Modifications and updates by: Jesse Smith
16  * Date: May 4, 2011
17  * Date: Jan 29, 2013
18  * Version 1.2 and newer
19  *
20  */
21
22
23 #include <getopt.h>
24 #include <stdio.h>
25 #include <stdlib.h>
26 #include <time.h>
27 #include <sys/time.h>
28 #include <sys/types.h>
29 #include <sys/wait.h>
30 #include <unistd.h>
31 #include <sys/types.h>
32 #include <signal.h>
33 #include <sys/resource.h>
34 #include <string.h>
35 #include <dirent.h>
36 #include <errno.h>
37 #include <string.h>
38 #include <limits.h>    // for compatibility
39
40 #ifdef __APPLE__
41 #include <mach/clock.h>
42 #include <mach/mach.h>
43 #endif
44
45
46 #ifdef FREEBSD
47 #include <fcntl.h>
48 #include <kvm.h>
49 #include <paths.h>
50 #include <sys/param.h>
51 #include <sys/sysctl.h>
52 #include <sys/user.h>
53 #endif
54
55
56 //kernel time resolution (inverse of one jiffy interval) in Hertz
57 //i don't know how to detect it, then define to the default (not very clean!)
58 #define HZ 100
59
60 //some useful macro
61 #define min(a,b) (a<b?a:b)
62 #define max(a,b) (a>b?a:b)
63
64 // For platforms without PATH_MAX
65 #ifndef PATH_MAX
66 #define PATH_MAX 4096
67 #endif
68
69 #define BEST_PRIORITY -10
70
71 #ifndef TRUE
72 #define TRUE 1
73 #endif
74 #ifndef FALSE
75 #define FALSE 0
76 #endif
77
78 #ifndef VERSION
79 #define VERSION 1.9
80 #endif
81
82 //pid of the controlled process
83 pid_t pid = 0;
84 pid_t my_pid;     // this process's PID
85
86 //executable file name
87 char *program_name;
88 //verbose mode
89 int verbose = FALSE;
90 //lazy mode
91 int lazy = FALSE;
92 // is higher priority nice possible?
93 int nice_lim;
94
95 // number of CPUs we detected
96 int NCPU;
97
98 //reverse byte search
99 // void *memrchr(const void *s, int c, size_t n);
100
101
102
103 //return ta-tb in microseconds (no overflow checks!)
104 inline long timediff(const struct timespec *ta,const struct timespec *tb) {
105     unsigned long us = (ta->tv_sec-tb->tv_sec)*1000000 + (ta->tv_nsec/1000 - tb->tv_nsec/1000);
106     return us;
107 }
108
109
110
111 int Check_Us(pid_t target_pid)
112 {
113    pid_t this_pid;
114
115    this_pid = getpid();
116    if (this_pid == target_pid)
117    {
118       printf("We cannot throttle ourselves.\n");
119       exit(7);
120    }
121    return TRUE;
122 }
123
124
125 int waitforpid(int pid) {
126         //switch to low priority
127         // if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
128         /*
129         if ( (nice_lim < INT_MAX) && 
130              (setpriority(PRIO_PROCESS, my_pid, 19) != 0) ) {
131                 printf("Warning: cannot renice\n");
132         }
133         */
134         int i=0;
135
136         while(1) {
137
138                 DIR *dip;
139                 struct dirent *dit;
140
141                 //open a directory stream to /proc directory
142                 if ((dip = opendir("/proc")) == NULL) {
143                         perror("opendir");
144                         return -1;
145                 }
146
147                 //read in from /proc and seek for process dirs
148                 while ((dit = readdir(dip)) != NULL) {
149                         //get pid
150                         if (pid==atoi(dit->d_name)) {
151                                 //pid detected
152                                 Check_Us(pid);
153                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
154                                         //process is ok!
155                                         if (closedir(dip) == -1) {
156                                            perror("closedir");
157                                            return -1;
158                                         }
159                                         goto done;
160                                 }
161                                 else {
162                                         fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
163                                 }
164                         }
165                 }
166
167                 //close the dir stream and check for errors
168                 if (closedir(dip) == -1) {
169                         perror("closedir");
170                         return -1;
171                 }
172
173                 //no suitable target found
174                 if (i++==0) {
175                         if (lazy) {
176                                 fprintf(stderr,"No process found\n");
177                                 exit(2);
178                         }
179                         else {
180                                 printf("Warning: no target process found. Waiting for it...\n");
181                         }
182                 }
183
184                 //sleep for a while
185                 sleep(2);
186         }
187
188 done:
189         printf("Process %d detected\n",pid);
190         //now set high priority, if possible
191         // if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
192         /*
193         if ( (nice_lim < INT_MAX) &&
194              (setpriority(PRIO_PROCESS, my_pid, nice_lim) != 0) ) {
195                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
196         }
197         */
198         return 0;
199
200 }
201
202 //this function periodically scans process list and looks for executable path names
203 //it should be executed in a low priority context, since precise timing does not matter
204 //if a process is found then its pid is returned
205 //process: the name of the wanted process, can be an absolute path name to the executable file
206 //         or simply its name
207 //return: pid of the found process
208 int getpidof(const char *process) {
209
210         //set low priority
211         // if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
212         /*
213         if ( (nice_lim < INT_MAX) &&
214              (setpriority(PRIO_PROCESS, my_pid, 19) != 0) ) {
215                 printf("Warning: cannot renice\n");
216         }
217         */
218         char exelink[20];
219         char exepath[PATH_MAX+1];
220         int pid=0;
221         int i=0;
222
223         while(1) {
224
225                 DIR *dip;
226                 struct dirent *dit;
227
228                 //open a directory stream to /proc directory
229                 if ((dip = opendir("/proc")) == NULL) {
230                         perror("opendir");
231                         return -1;
232                 }
233
234                 //read in from /proc and seek for process dirs
235                 while ((dit = readdir(dip)) != NULL) {
236                         //get pid
237                         pid=atoi(dit->d_name);
238                         if (pid>0) {
239                                 sprintf(exelink,"/proc/%d/exe",pid);
240                                 int size=readlink(exelink,exepath,sizeof(exepath));
241                                 if (size>0) {
242                                         int found=0;
243                                         if (process[0]=='/' && strncmp(exepath,process,size)==0 && size==strlen(process)) {
244                                                 //process starts with / then it's an absolute path
245                                                 found=1;
246                                         }
247                                         else {
248                                                 //process is the name of the executable file
249                                                 if (strncmp(exepath+size-strlen(process),process,strlen(process))==0) {
250                                                         found=1;
251                                                 }
252                                         }
253                                         if (found==1) {
254                                         Check_Us(pid);
255                                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
256                                                         //process is ok!
257                                                         if (closedir(dip) == -1) {
258                                                           perror("closedir");
259                                                           return -1;
260                                                         }
261                                                         goto done;
262                                                 }
263                                                 else {
264                                                         fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
265                                                 }
266                                         }
267                                 }
268                         }
269                 }
270
271                 //close the dir stream and check for errors
272                 if (closedir(dip) == -1) {
273                         perror("closedir");
274                         return -1;
275                 }
276
277                 //no suitable target found
278                 if (i++==0) {
279                         if (lazy) {
280                                 fprintf(stderr,"No process found\n");
281                                 exit(2);
282                         }
283                         else {
284                                 printf("Warning: no target process found. Waiting for it...\n");
285                         }
286                 }
287
288                 //sleep for a while
289                 sleep(2);
290         }
291
292 done:
293         printf("Process %d detected\n",pid);
294         //now set high priority, if possible
295         // if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
296         /*
297         if ( (nice_lim < INT_MAX) &&
298              (setpriority(PRIO_PROCESS, my_pid, nice_lim) != 0) ) {
299                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
300         }
301         */
302         return pid;
303
304 }
305
306 //SIGINT and SIGTERM signal handler
307 void quit(int sig) {
308         //let the process continue if it's stopped
309         kill(pid,SIGCONT);
310         printf("Exiting...\n");
311         exit(0);
312 }
313
314
315 #ifdef FREEBSD
316 //get jiffies count from /proc filesystem
317 int getjiffies(int pid)
318 {
319    kvm_t *my_kernel = NULL;
320    struct kinfo_proc *process_data = NULL;
321    int processes;
322    int my_jiffies = -1;
323
324    my_kernel = kvm_open(0, 0, 0, O_RDONLY, "kvm_open");
325    if (! my_kernel)
326    {
327       printf("Error opening kernel vm. You should be running as root.\n");
328       return -1;
329    }
330
331    process_data = kvm_getprocs(my_kernel, KERN_PROC_PID, pid, &processes);
332    if ( (process_data) && (processes >= 1) )
333        my_jiffies = process_data->ki_runtime;
334    
335    kvm_close(my_kernel);
336    if (my_jiffies >= 0)
337      my_jiffies /= 1000;
338    return my_jiffies;
339 }
340
341 #endif
342
343 #ifdef LINUX
344 int getjiffies(int pid) {
345         static char stat[20];
346         static char buffer[1024];
347         char *p;
348         sprintf(stat,"/proc/%d/stat",pid);
349         FILE *f=fopen(stat,"r");
350         if (f==NULL) return -1;
351         p = fgets(buffer,sizeof(buffer),f);
352         fclose(f);
353         // char *p=buffer;
354         if (p)
355         {
356           p=memchr(p+1,')',sizeof(buffer)-(p-buffer));
357           int sp=12;
358           while (sp--)
359                 p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
360           //user mode jiffies
361           int utime=atoi(p+1);
362           p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
363           //kernel mode jiffies
364           int ktime=atoi(p+1);
365           return utime+ktime;
366         }
367         // could not read info
368         return -1;
369 }
370 #endif
371
372 #ifdef __minix
373 int getjiffies(int pid)
374 {
375     char ps_info_path[32];
376     char buffer[1024];
377     char *p;
378     FILE *my_file;
379     int user_time, kernel_time;
380
381     sprintf(ps_info_path, "/proc/%d/psinfo", pid);
382     my_file = fopen(ps_info_path, "r");
383     if (! my_file)
384        return -1;
385     p = fgets(buffer, sizeof(buffer), my_file);
386     fclose(my_file);
387     if (p)
388     {
389         int spaces = 0;
390         while ( (spaces < 7) && (p[0]) )
391         {
392            if (p[0] == ' ')
393               spaces++;
394             p++;
395         }
396         user_time = atoi(p);
397         while ( (spaces < 8) && (p[0]) )
398         {
399             if (p[0] == ' ')
400               spaces++;
401             p++;
402         }
403         kernel_time = atoi(p);
404         // printf("%d %d\n", user_time, kernel_time);
405         return user_time + kernel_time;
406     }   // end of got contents of psinfo file
407     return -1;
408 }
409 #endif
410
411
412 //process instant photo
413 struct process_screenshot {
414         struct timespec when;   //timestamp
415         int jiffies;    //jiffies count of the process
416         int cputime;    //microseconds of work from previous screenshot to current
417 };
418
419 //extracted process statistics
420 struct cpu_usage {
421         float pcpu;
422         float workingrate;
423 };
424
425 //this function is an autonomous dynamic system
426 //it works with static variables (state variables of the system), that keep memory of recent past
427 //its aim is to estimate the cpu usage of the process
428 //to work properly it should be called in a fixed periodic way
429 //perhaps i will put it in a separate thread...
430 int compute_cpu_usage(int pid,int last_working_quantum,struct cpu_usage *pusage) {
431         #define MEM_ORDER 10
432         //circular buffer containing last MEM_ORDER process screenshots
433         static struct process_screenshot ps[MEM_ORDER];
434         //the last screenshot recorded in the buffer
435         static int front=-1;
436         //the oldest screenshot recorded in the buffer
437         static int tail=0;
438         #ifdef __minix
439         struct timeval tv;
440         #endif
441
442         if (pusage==NULL) {
443                 //reinit static variables
444                 front=-1;
445                 tail=0;
446                 return 0;
447         }
448
449         //let's advance front index and save the screenshot
450         front=(front+1)%MEM_ORDER;
451         int j=getjiffies(pid);
452         if (j>=0) ps[front].jiffies=j;
453         else return -1; //error: pid does not exist
454
455         #if defined(__APPLE__)
456         // OS X does not have clock_gettime, use clock_get_time
457         clock_serv_t cclock;
458         mach_timespec_t mts;
459         host_get_clock_service(mach_host_self(), CALENDAR_CLOCK, &cclock);
460         clock_get_time(cclock, &mts);
461         mach_port_deallocate(mach_task_self(), cclock);
462         ps[front].when.tv_sec = mts.tv_sec;
463         ps[front].when.tv_nsec = mts.tv_nsec;
464         #endif
465         #if defined(__minix)
466         // MINIX does not like the GNU style, nor the realtime style.
467         // try to do the best we can with gettimeofday()
468         gettimeofday(&tv, NULL);
469         ps[front].when.tv_sec = tv.tv_sec;
470         ps[front].when.tv_nsec = tv.tv_usec;
471         #endif
472         #if defined(LINUX) || defined(FREEBSD)
473         // Linux and BSD can use real time
474         clock_gettime(CLOCK_REALTIME,&(ps[front].when));
475         ps[front].cputime=last_working_quantum;
476         #endif
477         //buffer actual size is: (front-tail+MEM_ORDER)%MEM_ORDER+1
478         int size=(front-tail+MEM_ORDER)%MEM_ORDER+1;
479
480         if (size==1) {
481                 //not enough samples taken (it's the first one!), return -1
482                 pusage->pcpu=-1;
483                 pusage->workingrate=1;
484                 return 0;
485         }
486         else {
487                 //now we can calculate cpu usage, interval dt and dtwork are expressed in microseconds
488                 long dt=timediff(&(ps[front].when),&(ps[tail].when));
489                 long dtwork=0;
490                 int i=(tail+1)%MEM_ORDER;
491                 int max=(front+1)%MEM_ORDER;
492                 do {
493                         dtwork+=ps[i].cputime;
494                         i=(i+1)%MEM_ORDER;
495                 } while (i!=max);
496                 int used=ps[front].jiffies-ps[tail].jiffies;
497                 float usage=(used*1000000.0/HZ)/dtwork;
498                 pusage->workingrate=1.0*dtwork/dt;
499                 pusage->pcpu=usage*pusage->workingrate;
500                 if (size==MEM_ORDER)
501                         tail=(tail+1)%MEM_ORDER;
502                 return 0;
503         }
504         #undef MEM_ORDER
505 }
506
507 void print_caption() {
508         printf("\n%%CPU\twork quantum\tsleep quantum\tactive rate\n");
509 }
510
511
512 void increase_priority()
513 {
514         //find the best available nice value
515         int old_priority = getpriority(PRIO_PROCESS, 0);
516         int priority = old_priority;
517         while ( (setpriority(PRIO_PROCESS, 0, priority-1) == 0) &&
518                 (priority > BEST_PRIORITY) )
519         {
520                 priority--;
521         }
522         if (priority != old_priority) {
523                 if (verbose) printf("Priority changed to %d\n", priority);
524         }
525         else {
526                 if (verbose) printf("Warning: Cannot change priority. Run as root or renice for best results.\n");
527         }
528
529
530 }
531
532
533
534 void print_usage(FILE *stream,int exit_code) {
535         fprintf(stream, "CPUlimit version %1.1f\n", VERSION);
536         fprintf(stream, "Usage: %s TARGET [OPTIONS...] [PROGRAM]\n",program_name);
537         fprintf(stream, "   TARGET must be exactly one of these:\n");
538         fprintf(stream, "      -p, --pid=N        pid of the process\n");
539         fprintf(stream, "      -e, --exe=FILE     name of the executable program file\n");
540         fprintf(stream, "                         The -e option only works when\n");
541         fprintf(stream, "                         cpulimit is run with admin rights.\n");
542         fprintf(stream, "      -P, --path=PATH    absolute path name of the\n");
543         fprintf(stream, "                         executable program file\n");
544         fprintf(stream, "   OPTIONS\n");
545         fprintf(stream, "      -b  --background   run in background\n");
546         fprintf(stream, "      -c  --cpu=N        override the detection of CPUs on the machine.\n");
547         fprintf(stream, "      -l, --limit=N      percentage of cpu allowed from 1 up.\n");
548         fprintf(stream, "                         Usually 1 - %d00, but can be higher\n", NCPU);
549         fprintf(stream, "                         on multi-core CPUs (mandatory)\n");
550         fprintf(stream, "      -k, --kill         kill processes going over their limit\n");
551         fprintf(stream, "                         instead of just throttling them.\n");
552         fprintf(stream, "      -r, --restore      Restore processes after they have\n");
553         fprintf(stream, "                         been killed. Works with the -k flag.\n");
554         fprintf(stream, "      -v, --verbose      show control statistics\n");
555         fprintf(stream, "      -z, --lazy         exit if there is no suitable target process,\n");
556         fprintf(stream, "                         or if it dies\n");
557         fprintf(stream, "      -h, --help         display this help and exit\n");
558         exit(exit_code);
559 }
560
561
562
563 // Get the number of CPU cores on this machine.
564 int get_ncpu()
565 {
566         int ncpu = 1;
567 #ifdef _SC_NPROCESSORS_ONLN
568         ncpu = sysconf(_SC_NPROCESSORS_ONLN);
569 #endif
570         return ncpu;
571 }
572
573
574
575
576 int main(int argc, char **argv) {
577
578         //get program name
579         // char *p=(char*)memrchr(argv[0],(unsigned int)'/',strlen(argv[0]));
580         // program_name = p==NULL?argv[0]:(p+1);
581         program_name = argv[0];
582         int run_in_background = FALSE;
583         //parse arguments
584         int next_option;
585         /* A string listing valid short options letters. */
586         const char* short_options="p:e:P:l:c:bkrvzh";
587         /* An array describing valid long options. */
588         const struct option long_options[] = {
589                 { "pid", required_argument, NULL, 'p' },
590                 { "exe", required_argument, NULL, 'e' },
591                 { "path", required_argument, NULL, 'P' },
592                 { "limit", required_argument, NULL, 'l' },
593                 { "background", no_argument, NULL, 'b' },
594                 { "kill", no_argument, NULL, 'k' },
595                 { "restore", no_argument, NULL, 'r' },
596                 { "verbose", no_argument, NULL, 'v' },
597                 { "lazy", no_argument, NULL, 'z' },
598                 { "help", no_argument, NULL, 'h' },
599                 { "cpu", required_argument, NULL, 'c'},
600                 { NULL, 0, NULL, 0 }
601         };
602         //argument variables
603         const char *exe=NULL;
604         const char *path=NULL;
605         int perclimit=0;
606         int pid_ok = FALSE;
607         int process_ok = FALSE;
608         int limit_ok = FALSE;
609         int last_known_argument = 0;
610         int kill_process = FALSE;  // kill process instead of stopping it
611         int restore_process = FALSE;  // restore killed process
612         // struct rlimit maxlimit;
613
614         NCPU = get_ncpu();
615
616         do {
617                 next_option = getopt_long (argc, argv, short_options,long_options, NULL);
618                 switch(next_option) {
619                         case 'b':
620                                 run_in_background = TRUE;
621                                 last_known_argument++;
622                                 break;
623                         case 'p':
624                                 pid=atoi(optarg);
625                                 pid_ok = TRUE;
626                                 lazy = TRUE;
627                                 last_known_argument += 2;
628                                 break;
629                         case 'e':
630                                 exe=optarg;
631                                 process_ok = TRUE;
632                                 last_known_argument += 2;
633                                 break;
634                         case 'P':
635                                 path=optarg;
636                                 process_ok = TRUE;
637                                 last_known_argument += 2;
638                                 break;
639                         case 'l':
640                                 perclimit=atoi(optarg);
641                                 limit_ok = TRUE;
642                                 last_known_argument += 2;
643                                 break;
644                         case 'c':
645                                 NCPU = atoi(optarg);
646                                 last_known_argument += 2;
647                                 break;
648                         case 'k':
649                                 kill_process = TRUE;
650                                 last_known_argument++;
651                                 break;
652                         case 'r':
653                                 restore_process = TRUE;
654                                 last_known_argument++;
655                                 break;
656                         case 'v':
657                                 verbose = TRUE;
658                                 last_known_argument++;
659                                 break;
660                         case 'z':
661                                 lazy = TRUE;
662                                 last_known_argument++;
663                                 break;
664                         case 'h':
665                                 print_usage (stdout, 1);
666                                 last_known_argument++;
667                                 break;
668                         case '?':
669                                 print_usage (stderr, 1);
670                                 last_known_argument++;
671                                 break;
672                         case -1:
673                                 break;
674                         // default:
675                         //      abort();
676                 }
677         } while(next_option != -1);
678
679
680         if (last_known_argument + 1 < argc)
681         {
682            last_known_argument++;
683            pid_t forked_pid;
684            // try to launch remaining arguments
685            if (verbose)
686            {
687                int index = last_known_argument;
688                printf("Launching %s", argv[index]);
689                for (index = last_known_argument + 1; index < argc; index++)
690                     printf(" %s", argv[index]);
691                printf(" with limit %d\n", perclimit);
692            }
693            forked_pid = fork();
694            if (forked_pid == -1)  // error
695            {
696                printf("Failed to launch specified process.\n");
697                exit(1);
698            }
699            else if (forked_pid == 0)   // target child
700            {
701               execvp(argv[last_known_argument],
702                      &(argv[last_known_argument]) );
703               exit(2);
704            }
705            else     // parent who will now fork the throttler
706            {
707               pid_t limit_pid;
708               // if we are planning to kill a process, give it
709               // a running head start to avoid death at start-up
710               if (kill_process)
711                  sleep(5);
712               
713               limit_pid = fork();
714               if (limit_pid == 0)   // child
715               {
716                  pid = forked_pid;    // the first child
717                  lazy = TRUE;
718                  pid_ok = TRUE;
719                  if (verbose)
720                    printf("Throttling process %d\n", (int) pid);
721               }
722               else    // parent
723                 exit(0);
724            }
725
726            /*
727            else if (forked_pid == 0)   // child
728            {
729                lazy = TRUE;
730                pid_ok = TRUE;
731                pid = getppid();
732                if (verbose)
733                   printf("Throttling process %d\n", (int) pid);
734            }
735            else // parent
736            {
737                execvp(argv[last_known_argument], 
738                       &(argv[last_known_argument]));
739                
740                // we should never return  
741                exit(2);
742            }
743            */
744            
745         }      // end of launching child process
746
747         if (!process_ok && !pid_ok) {
748                 fprintf(stderr,"Error: You must specify a target process\n");
749                 print_usage (stderr, 1);
750                 exit(1);
751         }
752         if ((exe!=NULL && path!=NULL) || (pid_ok && (exe!=NULL || path!=NULL))) {
753                 fprintf(stderr,"Error: You must specify exactly one target process\n");
754                 print_usage (stderr, 1);
755                 exit(1);
756         }
757         if (!limit_ok) {
758                 fprintf(stderr,"Error: You must specify a cpu limit\n");
759                 print_usage (stderr, 1);
760                 exit(1);
761         }
762         float limit=perclimit/100.0;
763         if ( (limit <= 0.00) || (limit > NCPU) )
764         {
765                 fprintf(stderr,"Error: limit must be in the range of 1 to %d00\n", NCPU);
766                 print_usage (stderr, 1);
767                 exit(1);
768         }
769
770         // check to see if we should fork
771         if (run_in_background)
772         {
773              pid_t process_id;
774              process_id = fork();
775              if (! process_id)
776                 exit(0);
777              else
778              {
779                 setsid();
780                 process_id = fork();
781                 if (process_id)
782                   exit(0);
783              }
784         }
785
786         //parameters are all ok!
787         signal(SIGINT,quit);
788         signal(SIGTERM,quit);
789
790         my_pid = getpid();
791         if (verbose)
792            printf("%d CPUs detected.\n", NCPU);
793
794         increase_priority();
795 /*
796 Instead of all this big block of code to detect and change
797 priority settings, let us just use the increase_priority()
798 function. It is a little more simple and takes a more
799 gradual approach, rather than "all or nothing".
800 -- Jesse
801
802         if (setpriority(PRIO_PROCESS, my_pid,-20)!=0) {
803         //if that failed, check if we have a limit 
804         // by how much we can raise the priority
805 #ifdef RLIMIT_NICE 
806 //check if non-root can even make changes 
807 // (ifdef because it's only available in linux >= 2.6.13)
808                 nice_lim=getpriority(PRIO_PROCESS, my_pid);
809                 getrlimit(RLIMIT_NICE, &maxlimit);
810
811 //if we can do better then current
812                 if( (20 - (signed)maxlimit.rlim_cur) < nice_lim &&  
813                     setpriority(PRIO_PROCESS, my_pid,
814                     20 - (signed)maxlimit.rlim_cur)==0 //and it actually works
815                   ) {
816
817                         //if we can do better, but not by much, warn about it
818                         if( (nice_lim - (20 - (signed)maxlimit.rlim_cur)) < 9) 
819                         {
820                         printf("Warning, can only increase priority by %d.\n",                                nice_lim - (20 - (signed)maxlimit.rlim_cur));
821                         }
822                         //our new limit
823                         nice_lim = 20 - (signed)maxlimit.rlim_cur; 
824
825                 } else 
826 // otherwise don't try to change priority. 
827 // The below will also run if it's not possible 
828 // for non-root to change priority
829 #endif
830                 {
831                         printf("Warning: cannot renice.\nTo work better you should run this program as root, or adjust RLIMIT_NICE.\nFor example in /etc/security/limits.conf add a line with: * - nice -10\n\n");
832                         nice_lim=INT_MAX;
833                 }
834         } else {
835                 nice_lim=-20;
836         }
837 */
838
839
840         //time quantum in microseconds. it's splitted in a working period and a sleeping one
841         int period=100000;
842         struct timespec twork,tsleep;   //working and sleeping intervals
843         memset(&twork,0,sizeof(struct timespec));
844         memset(&tsleep,0,sizeof(struct timespec));
845
846 wait_for_process:
847
848         //look for the target process..or wait for it
849         if (exe != NULL)
850                 pid=getpidof(exe);
851         else if (path != NULL)
852                 pid=getpidof(path);
853         else 
854                 waitforpid(pid);
855
856         
857         //process detected...let's play
858
859         //init compute_cpu_usage internal stuff
860         compute_cpu_usage(0,0,NULL);
861         //main loop counter
862         int i=0;
863
864         struct timespec startwork,endwork;
865         long workingtime=0;             //last working time in microseconds
866
867         if (verbose) print_caption();
868
869         float pcpu_avg=0;
870
871         //here we should already have high priority, for time precision
872         while(1) {
873
874                 //estimate how much the controlled process is using the cpu in its working interval
875                 struct cpu_usage cu;
876                 if (compute_cpu_usage(pid,workingtime,&cu)==-1) {
877                         fprintf(stderr,"Process %d dead!\n",pid);
878                         if (lazy) exit(2);
879                         //wait until our process appears
880                         goto wait_for_process;          
881                 }
882
883                 //cpu actual usage of process (range 0-1)
884                 float pcpu=cu.pcpu;
885                 //rate at which we are keeping active the process (range 0-1)
886                 float workingrate=cu.workingrate;
887
888                 //adjust work and sleep time slices
889                 if (pcpu>0) {
890                         twork.tv_nsec=min(period*limit*1000/pcpu*workingrate,period*1000);
891                 }
892                 else if (pcpu==0) {
893                         twork.tv_nsec=period*1000;
894                 }
895                 else if (pcpu==-1) {
896                         //not yet a valid idea of cpu usage
897                         pcpu=limit;
898                         workingrate=limit;
899                         twork.tv_nsec=min(period*limit*1000,period*1000);
900                 }
901                 tsleep.tv_nsec=period*1000-twork.tv_nsec;
902
903                 //update average usage
904                 pcpu_avg=(pcpu_avg*i+pcpu)/(i+1);
905
906                 if (verbose && i%10==0 && i>0) {
907                         printf("%0.2f%%\t%6ld us\t%6ld us\t%0.2f%%\n",pcpu*100,twork.tv_nsec/1000,tsleep.tv_nsec/1000,workingrate*100);
908                         if (i%200 == 0)
909                            print_caption();
910                 }
911
912                 // if (limit<1 && limit>0) {
913                 // printf("Comparing %f to %f\n", pcpu, limit);
914                 if (pcpu < limit)
915                 {
916                         // printf("Continue\n");
917                         //resume process
918                         if (kill(pid,SIGCONT)!=0) {
919                                 fprintf(stderr,"Process %d dead!\n",pid);
920                                 if (lazy) exit(2);
921                                 //wait until our process appears
922                                 goto wait_for_process;
923                         }
924                 }
925
926                 #ifdef __APPLE__ 
927                 // OS X does not have clock_gettime, use clock_get_time
928                 clock_serv_t cclock;
929                 mach_timespec_t mts;
930                 host_get_clock_service(mach_host_self(), CALENDAR_CLOCK, &cclock);
931                 clock_get_time(cclock, &mts);
932                 mach_port_deallocate(mach_task_self(), cclock);
933                 startwork.tv_sec = mts.tv_sec;
934                 startwork.tv_nsec = mts.tv_nsec;
935                 #endif
936                 #ifdef __minix
937                 struct timeval my_tv;
938                 gettimeofday(&my_tv, NULL);
939                 startwork.tv_sec = my_tv.tv_sec;
940                 startwork.tv_nsec = my_tv.tv_usec;
941                 #endif
942                 #if defined(LINUX) || defined(FREEBSD)
943                 clock_gettime(CLOCK_REALTIME,&startwork);
944                 #endif
945
946                 nanosleep(&twork,NULL);         //now process is working
947                 #ifdef __APPLE__ 
948                 // OS X does not have clock_gettime, use clock_get_time
949                 // clock_serv_t cclock;
950                 // mach_timespec_t mts;
951                 host_get_clock_service(mach_host_self(), CALENDAR_CLOCK, &cclock);
952                 clock_get_time(cclock, &mts);
953                 mach_port_deallocate(mach_task_self(), cclock);
954                 endwork.tv_sec = mts.tv_sec;
955                 endwork.tv_nsec = mts.tv_nsec;
956                 #endif
957                 #ifdef __minix
958                 gettimeofday(&my_tv, NULL);
959                 endwork.tv_sec = my_tv.tv_sec;
960                 endwork.tv_nsec = my_tv.tv_usec;
961                 #endif
962                 #if defined(LINIX) || defined(FREEBSD)
963                 clock_gettime(CLOCK_REALTIME,&endwork);
964                 #endif
965                 workingtime=timediff(&endwork,&startwork);
966
967                 // if (limit<1) {
968                 // printf("Checking %f vs %f\n", pcpu, limit);
969                 if (pcpu > limit)
970                 {
971                      // When over our limit we may run into
972                      // situations where we want to kill
973                      // the offending process, then restart it
974                      if (kill_process)
975                      {
976                          kill(pid, SIGKILL);
977                          fprintf(stderr, "Process %d killed.\n", pid);
978                          if ( (lazy) && (! restore_process) ) 
979                               exit(2);
980                          // restart killed process
981                          if (restore_process)
982                          {
983                              pid_t new_process;
984                              new_process = fork();
985                              if (new_process == -1)
986                              {
987                               fprintf(stderr, "Failed to restore killed process.\n");
988                              }
989                              else if (new_process == 0)
990                              {
991                                 // child which becomes new process
992                                 if (verbose)
993                                    printf("Relaunching %s\n",
994                                           argv[last_known_argument]);
995                                 execvp(argv[last_known_argument],
996                                        &(argv[last_known_argument]) ); 
997                              }
998                              else // parent
999                              {
1000                                 // we need to track new process
1001                                 pid = new_process;
1002                                 // avoid killing child process
1003                                 sleep(5);
1004                              }
1005                          }
1006                      }
1007                      // do not kll process, just throttle it
1008                      else
1009                      {
1010                         // printf("Stop\n");
1011                         //stop process, it has worked enough
1012                         if (kill(pid,SIGSTOP)!=0) {
1013                                 fprintf(stderr,"Process %d dead!\n", pid);
1014                                 if (lazy) exit(2);
1015                                 //wait until our process appears
1016                                 goto wait_for_process;
1017                         }
1018                         nanosleep(&tsleep,NULL);        //now process is sleeping
1019                      }     // end of throttle process
1020                 }    // end of process using too much CPU
1021                 i++;
1022         }
1023
1024    return 0;
1025 }