Imported Upstream version 1.6
[debian/cpulimit.git] / cpulimit.c
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  * Version 1.2
18  *
19  */
20
21
22 #include <getopt.h>
23 #include <stdio.h>
24 #include <stdlib.h>
25 #include <time.h>
26 #include <sys/time.h>
27 #include <unistd.h>
28 #include <sys/types.h>
29 #include <signal.h>
30 #include <sys/resource.h>
31 #include <string.h>
32 #include <dirent.h>
33 #include <errno.h>
34 #include <string.h>
35 #include <limits.h>    // for compatibility
36
37
38 //kernel time resolution (inverse of one jiffy interval) in Hertz
39 //i don't know how to detect it, then define to the default (not very clean!)
40 #define HZ 100
41
42 //some useful macro
43 #define min(a,b) (a<b?a:b)
44 #define max(a,b) (a>b?a:b)
45
46 // For platforms without PATH_MAX
47 #ifndef PATH_MAX
48 #define PATH_MAX 4096
49 #endif
50
51 //pid of the controlled process
52 int pid=0;
53 //executable file name
54 char *program_name;
55 //verbose mode
56 int verbose=0;
57 //lazy mode
58 int lazy=0;
59 // is higher priority nice possible?
60 int nice_lim;
61
62 //reverse byte search
63 void *memrchr(const void *s, int c, size_t n);
64
65 //return ta-tb in microseconds (no overflow checks!)
66 inline long timediff(const struct timespec *ta,const struct timespec *tb) {
67     unsigned long us = (ta->tv_sec-tb->tv_sec)*1000000 + (ta->tv_nsec/1000 - tb->tv_nsec/1000);
68     return us;
69 }
70
71 int waitforpid(int pid) {
72         //switch to low priority
73         // if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
74         if ( (nice_lim < INT_MAX) && 
75              (setpriority(PRIO_PROCESS, getpid(), 19) != 0) ) {
76                 printf("Warning: cannot renice\n");
77         }
78
79         int i=0;
80
81         while(1) {
82
83                 DIR *dip;
84                 struct dirent *dit;
85
86                 //open a directory stream to /proc directory
87                 if ((dip = opendir("/proc")) == NULL) {
88                         perror("opendir");
89                         return -1;
90                 }
91
92                 //read in from /proc and seek for process dirs
93                 while ((dit = readdir(dip)) != NULL) {
94                         //get pid
95                         if (pid==atoi(dit->d_name)) {
96                                 //pid detected
97                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
98                                         //process is ok!
99                                         if (closedir(dip) == -1) {
100                                            perror("closedir");
101                                            return -1;
102                                         }
103                                         goto done;
104                                 }
105                                 else {
106                                         fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
107                                 }
108                         }
109                 }
110
111                 //close the dir stream and check for errors
112                 if (closedir(dip) == -1) {
113                         perror("closedir");
114                         return -1;
115                 }
116
117                 //no suitable target found
118                 if (i++==0) {
119                         if (lazy) {
120                                 fprintf(stderr,"No process found\n");
121                                 exit(2);
122                         }
123                         else {
124                                 printf("Warning: no target process found. Waiting for it...\n");
125                         }
126                 }
127
128                 //sleep for a while
129                 sleep(2);
130         }
131
132 done:
133         printf("Process %d detected\n",pid);
134         //now set high priority, if possible
135         // if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
136         if ( (nice_lim < INT_MAX) &&
137              (setpriority(PRIO_PROCESS, getpid(), nice_lim) != 0) ) {
138                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
139         }
140         return 0;
141
142 }
143
144 //this function periodically scans process list and looks for executable path names
145 //it should be executed in a low priority context, since precise timing does not matter
146 //if a process is found then its pid is returned
147 //process: the name of the wanted process, can be an absolute path name to the executable file
148 //         or simply its name
149 //return: pid of the found process
150 int getpidof(const char *process) {
151
152         //set low priority
153         // if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
154         if ( (nice_lim < INT_MAX) &&
155              (setpriority(PRIO_PROCESS, getpid(), 19) != 0) ) {
156                 printf("Warning: cannot renice\n");
157         }
158
159         char exelink[20];
160         char exepath[PATH_MAX+1];
161         int pid=0;
162         int i=0;
163
164         while(1) {
165
166                 DIR *dip;
167                 struct dirent *dit;
168
169                 //open a directory stream to /proc directory
170                 if ((dip = opendir("/proc")) == NULL) {
171                         perror("opendir");
172                         return -1;
173                 }
174
175                 //read in from /proc and seek for process dirs
176                 while ((dit = readdir(dip)) != NULL) {
177                         //get pid
178                         pid=atoi(dit->d_name);
179                         if (pid>0) {
180                                 sprintf(exelink,"/proc/%d/exe",pid);
181                                 int size=readlink(exelink,exepath,sizeof(exepath));
182                                 if (size>0) {
183                                         int found=0;
184                                         if (process[0]=='/' && strncmp(exepath,process,size)==0 && size==strlen(process)) {
185                                                 //process starts with / then it's an absolute path
186                                                 found=1;
187                                         }
188                                         else {
189                                                 //process is the name of the executable file
190                                                 if (strncmp(exepath+size-strlen(process),process,strlen(process))==0) {
191                                                         found=1;
192                                                 }
193                                         }
194                                         if (found==1) {
195                                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
196                                                         //process is ok!
197                                                         if (closedir(dip) == -1) {
198                                                           perror("closedir");
199                                                           return -1;
200                                                         }
201                                                         goto done;
202                                                 }
203                                                 else {
204                                                         fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
205                                                 }
206                                         }
207                                 }
208                         }
209                 }
210
211                 //close the dir stream and check for errors
212                 if (closedir(dip) == -1) {
213                         perror("closedir");
214                         return -1;
215                 }
216
217                 //no suitable target found
218                 if (i++==0) {
219                         if (lazy) {
220                                 fprintf(stderr,"No process found\n");
221                                 exit(2);
222                         }
223                         else {
224                                 printf("Warning: no target process found. Waiting for it...\n");
225                         }
226                 }
227
228                 //sleep for a while
229                 sleep(2);
230         }
231
232 done:
233         printf("Process %d detected\n",pid);
234         //now set high priority, if possible
235         // if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
236         if ( (nice_lim < INT_MAX) &&
237              (setpriority(PRIO_PROCESS, getpid(), nice_lim) != 0) ) {
238                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
239         }
240         return pid;
241
242 }
243
244 //SIGINT and SIGTERM signal handler
245 void quit(int sig) {
246         //let the process continue if it's stopped
247         kill(pid,SIGCONT);
248         printf("Exiting...\n");
249         exit(0);
250 }
251
252 //get jiffies count from /proc filesystem
253 int getjiffies(int pid) {
254         static char stat[20];
255         static char buffer[1024];
256         char *p;
257         sprintf(stat,"/proc/%d/stat",pid);
258         FILE *f=fopen(stat,"r");
259         if (f==NULL) return -1;
260         p = fgets(buffer,sizeof(buffer),f);
261         fclose(f);
262         // char *p=buffer;
263         if (p)
264         {
265           p=memchr(p+1,')',sizeof(buffer)-(p-buffer));
266           int sp=12;
267           while (sp--)
268                 p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
269           //user mode jiffies
270           int utime=atoi(p+1);
271           p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
272           //kernel mode jiffies
273           int ktime=atoi(p+1);
274           return utime+ktime;
275         }
276         // could not read info
277         return -1;
278 }
279
280 //process instant photo
281 struct process_screenshot {
282         struct timespec when;   //timestamp
283         int jiffies;    //jiffies count of the process
284         int cputime;    //microseconds of work from previous screenshot to current
285 };
286
287 //extracted process statistics
288 struct cpu_usage {
289         float pcpu;
290         float workingrate;
291 };
292
293 //this function is an autonomous dynamic system
294 //it works with static variables (state variables of the system), that keep memory of recent past
295 //its aim is to estimate the cpu usage of the process
296 //to work properly it should be called in a fixed periodic way
297 //perhaps i will put it in a separate thread...
298 int compute_cpu_usage(int pid,int last_working_quantum,struct cpu_usage *pusage) {
299         #define MEM_ORDER 10
300         //circular buffer containing last MEM_ORDER process screenshots
301         static struct process_screenshot ps[MEM_ORDER];
302         //the last screenshot recorded in the buffer
303         static int front=-1;
304         //the oldest screenshot recorded in the buffer
305         static int tail=0;
306
307         if (pusage==NULL) {
308                 //reinit static variables
309                 front=-1;
310                 tail=0;
311                 return 0;
312         }
313
314         //let's advance front index and save the screenshot
315         front=(front+1)%MEM_ORDER;
316         int j=getjiffies(pid);
317         if (j>=0) ps[front].jiffies=j;
318         else return -1; //error: pid does not exist
319         clock_gettime(CLOCK_REALTIME,&(ps[front].when));
320         ps[front].cputime=last_working_quantum;
321
322         //buffer actual size is: (front-tail+MEM_ORDER)%MEM_ORDER+1
323         int size=(front-tail+MEM_ORDER)%MEM_ORDER+1;
324
325         if (size==1) {
326                 //not enough samples taken (it's the first one!), return -1
327                 pusage->pcpu=-1;
328                 pusage->workingrate=1;
329                 return 0;
330         }
331         else {
332                 //now we can calculate cpu usage, interval dt and dtwork are expressed in microseconds
333                 long dt=timediff(&(ps[front].when),&(ps[tail].when));
334                 long dtwork=0;
335                 int i=(tail+1)%MEM_ORDER;
336                 int max=(front+1)%MEM_ORDER;
337                 do {
338                         dtwork+=ps[i].cputime;
339                         i=(i+1)%MEM_ORDER;
340                 } while (i!=max);
341                 int used=ps[front].jiffies-ps[tail].jiffies;
342                 float usage=(used*1000000.0/HZ)/dtwork;
343                 pusage->workingrate=1.0*dtwork/dt;
344                 pusage->pcpu=usage*pusage->workingrate;
345                 if (size==MEM_ORDER)
346                         tail=(tail+1)%MEM_ORDER;
347                 return 0;
348         }
349         #undef MEM_ORDER
350 }
351
352 void print_caption() {
353         printf("\n%%CPU\twork quantum\tsleep quantum\tactive rate\n");
354 }
355
356 void print_usage(FILE *stream,int exit_code) {
357         fprintf(stream, "Usage: %s TARGET [OPTIONS...]\n",program_name);
358         fprintf(stream, "   TARGET must be exactly one of these:\n");
359         fprintf(stream, "      -p, --pid=N        pid of the process\n");
360         fprintf(stream, "      -e, --exe=FILE     name of the executable program file\n");
361         fprintf(stream, "                         The -e option only works when\n");
362         fprintf(stream, "                         cpulimit is run with admin rights.\n");
363         fprintf(stream, "      -P, --path=PATH    absolute path name of the\n");
364         fprintf(stream, "                         executable program file\n");
365         fprintf(stream, "   OPTIONS\n");
366         fprintf(stream, "      -b  --background   run in background\n");
367         fprintf(stream, "      -l, --limit=N      percentage of cpu allowed from 1 up.\n");
368         fprintf(stream, "                         Usually 1 - 100, but can be higher\n");
369         fprintf(stream, "                         on multi-core CPUs (mandatory)\n");
370         fprintf(stream, "      -v, --verbose      show control statistics\n");
371         fprintf(stream, "      -z, --lazy         exit if there is no suitable target process,\n");
372         fprintf(stream, "                         or if it dies\n");
373         fprintf(stream, "      -h, --help         display this help and exit\n");
374         exit(exit_code);
375 }
376
377 int main(int argc, char **argv) {
378
379         //get program name
380         char *p=(char*)memrchr(argv[0],(unsigned int)'/',strlen(argv[0]));
381         program_name = p==NULL?argv[0]:(p+1);
382         int run_in_background = 0;
383         //parse arguments
384         int next_option;
385         /* A string listing valid short options letters. */
386         const char* short_options="p:e:P:l:bvzh";
387         /* An array describing valid long options. */
388         const struct option long_options[] = {
389                 { "pid", required_argument, NULL, 'p' },
390                 { "exe", required_argument, NULL, 'e' },
391                 { "path", required_argument, NULL, 'P' },
392                 { "limit", required_argument, NULL, 'l' },
393                 { "background", no_argument, NULL, 'b' },
394                 { "verbose", no_argument, NULL, 'v' },
395                 { "lazy", no_argument, NULL, 'z' },
396                 { "help", no_argument, NULL, 'h' },
397                 { NULL, 0, NULL, 0 }
398         };
399         //argument variables
400         const char *exe=NULL;
401         const char *path=NULL;
402         int perclimit=0;
403         int pid_ok=0;
404         int process_ok=0;
405         int limit_ok=0;
406         struct rlimit maxlimit;
407
408         do {
409                 next_option = getopt_long (argc, argv, short_options,long_options, NULL);
410                 switch(next_option) {
411                         case 'b':
412                                 run_in_background = 1;
413                                 break;
414                         case 'p':
415                                 pid=atoi(optarg);
416                                 pid_ok=1;
417                                 lazy = 1;
418                                 break;
419                         case 'e':
420                                 exe=optarg;
421                                 process_ok=1;
422                                 break;
423                         case 'P':
424                                 path=optarg;
425                                 process_ok=1;
426                                 break;
427                         case 'l':
428                                 perclimit=atoi(optarg);
429                                 limit_ok=1;
430                                 break;
431                         case 'v':
432                                 verbose=1;
433                                 break;
434                         case 'z':
435                                 lazy=1;
436                                 break;
437                         case 'h':
438                                 print_usage (stdout, 1);
439                                 break;
440                         case '?':
441                                 print_usage (stderr, 1);
442                                 break;
443                         case -1:
444                                 break;
445                         default:
446                                 abort();
447                 }
448         } while(next_option != -1);
449
450         if (!process_ok && !pid_ok) {
451                 fprintf(stderr,"Error: You must specify a target process\n");
452                 print_usage (stderr, 1);
453                 exit(1);
454         }
455         if ((exe!=NULL && path!=NULL) || (pid_ok && (exe!=NULL || path!=NULL))) {
456                 fprintf(stderr,"Error: You must specify exactly one target process\n");
457                 print_usage (stderr, 1);
458                 exit(1);
459         }
460         if (!limit_ok) {
461                 fprintf(stderr,"Error: You must specify a cpu limit\n");
462                 print_usage (stderr, 1);
463                 exit(1);
464         }
465         float limit=perclimit/100.0;
466         if (limit <= 0.00) // || limit >1) {
467         {
468                 fprintf(stderr,"Error: limit must be greater than 0\n");
469                 print_usage (stderr, 1);
470                 exit(1);
471         }
472
473         // check to see if we should fork
474         if (run_in_background)
475         {
476              pid_t process_id;
477              process_id = fork();
478              if (! process_id)
479                 exit(0);
480              else
481              {
482                 setsid();
483                 process_id = fork();
484                 if (process_id)
485                   exit(0);
486              }
487         }
488
489         //parameters are all ok!
490         signal(SIGINT,quit);
491         signal(SIGTERM,quit);
492
493         if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
494         //if that failed, check if we have a limit 
495         // by how much we can raise the priority
496 #ifdef RLIMIT_NICE 
497 //check if non-root can even make changes 
498 // (ifdef because it's only available in linux >= 2.6.13)
499                 nice_lim=getpriority(PRIO_PROCESS,getpid());
500                 getrlimit(RLIMIT_NICE, &maxlimit);
501
502 //if we can do better then current
503                 if( (20 - (signed)maxlimit.rlim_cur) < nice_lim &&  
504                     setpriority(PRIO_PROCESS,getpid(),
505                     20 - (signed)maxlimit.rlim_cur)==0 //and it actually works
506                   ) {
507
508                         //if we can do better, but not by much, warn about it
509                         if( (nice_lim - (20 - (signed)maxlimit.rlim_cur)) < 9) 
510                         {
511                         printf("Warning, can only increase priority by %d.\n",                                nice_lim - (20 - (signed)maxlimit.rlim_cur));
512                         }
513                         //our new limit
514                         nice_lim = 20 - (signed)maxlimit.rlim_cur; 
515
516                 } else 
517 // otherwise don't try to change priority. 
518 // The below will also run if it's not possible 
519 // for non-root to change priority
520 #endif
521                 {
522                         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");
523                         nice_lim=INT_MAX;
524                 }
525         } else {
526                 nice_lim=-20;
527         }
528         //don't bother putting setpriority back down, 
529         // since getpidof and waitforpid twiddle it anyway
530
531
532
533         //time quantum in microseconds. it's splitted in a working period and a sleeping one
534         int period=100000;
535         struct timespec twork,tsleep;   //working and sleeping intervals
536         memset(&twork,0,sizeof(struct timespec));
537         memset(&tsleep,0,sizeof(struct timespec));
538
539 wait_for_process:
540
541         //look for the target process..or wait for it
542         if (exe!=NULL)
543                 pid=getpidof(exe);
544         else if (path!=NULL)
545                 pid=getpidof(path);
546         else {
547                 waitforpid(pid);
548         }
549         //process detected...let's play
550
551         //init compute_cpu_usage internal stuff
552         compute_cpu_usage(0,0,NULL);
553         //main loop counter
554         int i=0;
555
556         struct timespec startwork,endwork;
557         long workingtime=0;             //last working time in microseconds
558
559         if (verbose) print_caption();
560
561         float pcpu_avg=0;
562
563         //here we should already have high priority, for time precision
564         while(1) {
565
566                 //estimate how much the controlled process is using the cpu in its working interval
567                 struct cpu_usage cu;
568                 if (compute_cpu_usage(pid,workingtime,&cu)==-1) {
569                         fprintf(stderr,"Process %d dead!\n",pid);
570                         if (lazy) exit(2);
571                         //wait until our process appears
572                         goto wait_for_process;          
573                 }
574
575                 //cpu actual usage of process (range 0-1)
576                 float pcpu=cu.pcpu;
577                 //rate at which we are keeping active the process (range 0-1)
578                 float workingrate=cu.workingrate;
579
580                 //adjust work and sleep time slices
581                 if (pcpu>0) {
582                         twork.tv_nsec=min(period*limit*1000/pcpu*workingrate,period*1000);
583                 }
584                 else if (pcpu==0) {
585                         twork.tv_nsec=period*1000;
586                 }
587                 else if (pcpu==-1) {
588                         //not yet a valid idea of cpu usage
589                         pcpu=limit;
590                         workingrate=limit;
591                         twork.tv_nsec=min(period*limit*1000,period*1000);
592                 }
593                 tsleep.tv_nsec=period*1000-twork.tv_nsec;
594
595                 //update average usage
596                 pcpu_avg=(pcpu_avg*i+pcpu)/(i+1);
597
598                 if (verbose && i%10==0 && i>0) {
599                         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);
600                 }
601
602                 if (limit<1 && limit>0) {
603                         //resume process
604                         if (kill(pid,SIGCONT)!=0) {
605                                 fprintf(stderr,"Process %d dead!\n",pid);
606                                 if (lazy) exit(2);
607                                 //wait until our process appears
608                                 goto wait_for_process;
609                         }
610                 }
611
612                 clock_gettime(CLOCK_REALTIME,&startwork);
613                 nanosleep(&twork,NULL);         //now process is working        
614                 clock_gettime(CLOCK_REALTIME,&endwork);
615                 workingtime=timediff(&endwork,&startwork);
616
617                 if (limit<1) {
618                         //stop process, it has worked enough
619                         if (kill(pid,SIGSTOP)!=0) {
620                                 fprintf(stderr,"Process %d dead!\n",pid);
621                                 if (lazy) exit(2);
622                                 //wait until our process appears
623                                 goto wait_for_process;
624                         }
625                         nanosleep(&tsleep,NULL);        //now process is sleeping
626                 }
627                 i++;
628         }
629
630 }