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