Load cpulimit-1.1 into debian/cpulimit/trunk.
[debian/cpulimit.git] / cpulimit.c
1 /**
2  * Copyright (c) 2005, by:      Angelo Marletta <marlonx80@hotmail.com>
3  *
4  * This file may be used subject to the terms and conditions of the
5  * GNU Library General Public License Version 2, or any later version
6  * at your option, as published by the Free Software Foundation.
7  * This program is distributed in the hope that it will be useful,
8  * but WITHOUT ANY WARRANTY; without even the implied warranty of
9  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
10  * GNU Library General Public License for more details.
11  *
12  **********************************************************************
13  *
14  * Simple program to limit the cpu usage of a process
15  * If you modify this code, send me a copy please
16  *
17  * Author:  Angelo Marletta
18  * Date:    26/06/2005
19  * Version: 1.1
20  * Last version at: http://marlon80.interfree.it/cpulimit/index.html
21  *
22  * Changelog:
23  *  - Fixed a segmentation fault if controlled process exited in particular circumstances
24  *  - Better CPU usage estimate
25  *  - Fixed a <0 %CPU usage reporting in rare cases
26  *  - Replaced MAX_PATH_SIZE with PATH_MAX already defined in <limits.h>
27  *  - Command line arguments now available
28  *  - Now is possible to specify target process by pid
29  */
30
31
32 #include <getopt.h>
33 #include <stdio.h>
34 #include <stdlib.h>
35 #include <time.h>
36 #include <sys/time.h>
37 #include <unistd.h>
38 #include <sys/types.h>
39 #include <signal.h>
40 #include <sys/resource.h>
41 #include <string.h>
42 #include <dirent.h>
43 #include <errno.h>
44 #include <string.h>
45
46 //kernel time resolution (inverse of one jiffy interval) in Hertz
47 //i don't know how to detect it, then define to the default (not very clean!)
48 #define HZ 100
49
50 //some useful macro
51 #define min(a,b) (a<b?a:b)
52 #define max(a,b) (a>b?a:b)
53
54 //pid of the controlled process
55 int pid=0;
56 //executable file name
57 char *program_name;
58 //verbose mode
59 int verbose=0;
60 //lazy mode
61 int lazy=0;
62
63 //reverse byte search
64 void *memrchr(const void *s, int c, size_t n);
65
66 //return ta-tb in microseconds (no overflow checks!)
67 inline long timediff(const struct timespec *ta,const struct timespec *tb) {
68     unsigned long us = (ta->tv_sec-tb->tv_sec)*1000000 + (ta->tv_nsec/1000 - tb->tv_nsec/1000);
69     return us;
70 }
71
72 int waitforpid(int pid) {
73         //switch to low priority
74         if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
75                 printf("Warning: cannot renice\n");
76         }
77
78         int i=0;
79
80         while(1) {
81
82                 DIR *dip;
83                 struct dirent *dit;
84
85                 //open a directory stream to /proc directory
86                 if ((dip = opendir("/proc")) == NULL) {
87                         perror("opendir");
88                         return -1;
89                 }
90
91                 //read in from /proc and seek for process dirs
92                 while ((dit = readdir(dip)) != NULL) {
93                         //get pid
94                         if (pid==atoi(dit->d_name)) {
95                                 //pid detected
96                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
97                                         //process is ok!
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                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
132         }
133         return 0;
134
135 }
136
137 //this function periodically scans process list and looks for executable path names
138 //it should be executed in a low priority context, since precise timing does not matter
139 //if a process is found then its pid is returned
140 //process: the name of the wanted process, can be an absolute path name to the executable file
141 //         or simply its name
142 //return: pid of the found process
143 int getpidof(const char *process) {
144
145         //set low priority
146         if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
147                 printf("Warning: cannot renice\n");
148         }
149
150         char exelink[20];
151         char exepath[PATH_MAX+1];
152         int pid=0;
153         int i=0;
154
155         while(1) {
156
157                 DIR *dip;
158                 struct dirent *dit;
159
160                 //open a directory stream to /proc directory
161                 if ((dip = opendir("/proc")) == NULL) {
162                         perror("opendir");
163                         return -1;
164                 }
165
166                 //read in from /proc and seek for process dirs
167                 while ((dit = readdir(dip)) != NULL) {
168                         //get pid
169                         pid=atoi(dit->d_name);
170                         if (pid>0) {
171                                 sprintf(exelink,"/proc/%d/exe",pid);
172                                 int size=readlink(exelink,exepath,sizeof(exepath));
173                                 if (size>0) {
174                                         int found=0;
175                                         if (process[0]=='/' && strncmp(exepath,process,size)==0 && size==strlen(process)) {
176                                                 //process starts with / then it's an absolute path
177                                                 found=1;
178                                         }
179                                         else {
180                                                 //process is the name of the executable file
181                                                 if (strncmp(exepath+size-strlen(process),process,strlen(process))==0) {
182                                                         found=1;
183                                                 }
184                                         }
185                                         if (found==1) {
186                                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
187                                                         //process is ok!
188                                                         goto done;
189                                                 }
190                                                 else {
191                                                         fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
192                                                 }
193                                         }
194                                 }
195                         }
196                 }
197
198                 //close the dir stream and check for errors
199                 if (closedir(dip) == -1) {
200                         perror("closedir");
201                         return -1;
202                 }
203
204                 //no suitable target found
205                 if (i++==0) {
206                         if (lazy) {
207                                 fprintf(stderr,"No process found\n");
208                                 exit(2);
209                         }
210                         else {
211                                 printf("Warning: no target process found. Waiting for it...\n");
212                         }
213                 }
214
215                 //sleep for a while
216                 sleep(2);
217         }
218
219 done:
220         printf("Process %d detected\n",pid);
221         //now set high priority, if possible
222         if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
223                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
224         }
225         return pid;
226
227 }
228
229 //SIGINT and SIGTERM signal handler
230 void quit(int sig) {
231         //let the process continue if it's stopped
232         kill(pid,SIGCONT);
233         printf("Exiting...\n");
234         exit(0);
235 }
236
237 //get jiffies count from /proc filesystem
238 int getjiffies(int pid) {
239         static char stat[20];
240         static char buffer[1024];
241         sprintf(stat,"/proc/%d/stat",pid);
242         FILE *f=fopen(stat,"r");
243         if (f==NULL) return -1;
244         fgets(buffer,sizeof(buffer),f);
245         fclose(f);
246         char *p=buffer;
247         p=memchr(p+1,')',sizeof(buffer)-(p-buffer));
248         int sp=12;
249         while (sp--)
250                 p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
251         //user mode jiffies
252         int utime=atoi(p+1);
253         p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
254         //kernel mode jiffies
255         int ktime=atoi(p+1);
256         return utime+ktime;
257 }
258
259 //process instant photo
260 struct process_screenshot {
261         struct timespec when;   //timestamp
262         int jiffies;    //jiffies count of the process
263         int cputime;    //microseconds of work from previous screenshot to current
264 };
265
266 //extracted process statistics
267 struct cpu_usage {
268         float pcpu;
269         float workingrate;
270 };
271
272 //this function is an autonomous dynamic system
273 //it works with static variables (state variables of the system), that keep memory of recent past
274 //its aim is to estimate the cpu usage of the process
275 //to work properly it should be called in a fixed periodic way
276 //perhaps i will put it in a separate thread...
277 int compute_cpu_usage(int pid,int last_working_quantum,struct cpu_usage *pusage) {
278         #define MEM_ORDER 10
279         //circular buffer containing last MEM_ORDER process screenshots
280         static struct process_screenshot ps[MEM_ORDER];
281         //the last screenshot recorded in the buffer
282         static int front=-1;
283         //the oldest screenshot recorded in the buffer
284         static int tail=0;
285
286         if (pusage==NULL) {
287                 //reinit static variables
288                 front=-1;
289                 tail=0;
290                 return 0;
291         }
292
293         //let's advance front index and save the screenshot
294         front=(front+1)%MEM_ORDER;
295         int j=getjiffies(pid);
296         if (j>=0) ps[front].jiffies=j;
297         else return -1; //error: pid does not exist
298         clock_gettime(CLOCK_REALTIME,&(ps[front].when));
299         ps[front].cputime=last_working_quantum;
300
301         //buffer actual size is: (front-tail+MEM_ORDER)%MEM_ORDER+1
302         int size=(front-tail+MEM_ORDER)%MEM_ORDER+1;
303
304         if (size==1) {
305                 //not enough samples taken (it's the first one!), return -1
306                 pusage->pcpu=-1;
307                 pusage->workingrate=1;
308                 return 0;
309         }
310         else {
311                 //now we can calculate cpu usage, interval dt and dtwork are expressed in microseconds
312                 long dt=timediff(&(ps[front].when),&(ps[tail].when));
313                 long dtwork=0;
314                 int i=(tail+1)%MEM_ORDER;
315                 int max=(front+1)%MEM_ORDER;
316                 do {
317                         dtwork+=ps[i].cputime;
318                         i=(i+1)%MEM_ORDER;
319                 } while (i!=max);
320                 int used=ps[front].jiffies-ps[tail].jiffies;
321                 float usage=(used*1000000.0/HZ)/dtwork;
322                 pusage->workingrate=1.0*dtwork/dt;
323                 pusage->pcpu=usage*pusage->workingrate;
324                 if (size==MEM_ORDER)
325                         tail=(tail+1)%MEM_ORDER;
326                 return 0;
327         }
328         #undef MEM_ORDER
329 }
330
331 void print_caption() {
332         printf("\n%%CPU\twork quantum\tsleep quantum\tactive rate\n");
333 }
334
335 void print_usage(FILE *stream,int exit_code) {
336         fprintf(stream, "Usage: %s TARGET [OPTIONS...]\n",program_name);
337         fprintf(stream, "   TARGET must be exactly one of these:\n");
338         fprintf(stream, "      -p, --pid=N        pid of the process\n");
339         fprintf(stream, "      -e, --exe=FILE     name of the executable program file\n");
340         fprintf(stream, "      -P, --path=PATH    absolute path name of the executable program file\n");
341         fprintf(stream, "   OPTIONS\n");
342         fprintf(stream, "      -l, --limit=N      percentage of cpu allowed from 0 to 100 (mandatory)\n");
343         fprintf(stream, "      -v, --verbose      show control statistics\n");
344         fprintf(stream, "      -z, --lazy         exit if there is no suitable target process, or if it dies\n");
345         fprintf(stream, "      -h, --help         display this help and exit\n");
346         exit(exit_code);
347 }
348
349 int main(int argc, char **argv) {
350
351         //get program name
352         char *p=(char*)memrchr(argv[0],(unsigned int)'/',strlen(argv[0]));
353         program_name = p==NULL?argv[0]:(p+1);
354         //parse arguments
355         int next_option;
356         /* A string listing valid short options letters. */
357         const char* short_options="p:e:P:l:vzh";
358         /* An array describing valid long options. */
359         const struct option long_options[] = {
360                 { "pid", 0, NULL, 'p' },
361                 { "exe", 1, NULL, 'e' },
362                 { "path", 0, NULL, 'P' },
363                 { "limit", 0, NULL, 'l' },
364                 { "verbose", 0, NULL, 'v' },
365                 { "lazy", 0, NULL, 'z' },
366                 { "help", 0, NULL, 'h' },
367                 { NULL, 0, NULL, 0 }
368         };
369         //argument variables
370         const char *exe=NULL;
371         const char *path=NULL;
372         int perclimit=0;
373         int pid_ok=0;
374         int process_ok=0;
375         int limit_ok=0;
376
377         do {
378                 next_option = getopt_long (argc, argv, short_options,long_options, NULL);
379                 switch(next_option) {
380                         case 'p':
381                                 pid=atoi(optarg);
382                                 pid_ok=1;
383                                 break;
384                         case 'e':
385                                 exe=optarg;
386                                 process_ok=1;
387                                 break;
388                         case 'P':
389                                 path=optarg;
390                                 process_ok=1;
391                                 break;
392                         case 'l':
393                                 perclimit=atoi(optarg);
394                                 limit_ok=1;
395                                 break;
396                         case 'v':
397                                 verbose=1;
398                                 break;
399                         case 'z':
400                                 lazy=1;
401                                 break;
402                         case 'h':
403                                 print_usage (stdout, 1);
404                                 break;
405                         case '?':
406                                 print_usage (stderr, 1);
407                                 break;
408                         case -1:
409                                 break;
410                         default:
411                                 abort();
412                 }
413         } while(next_option != -1);
414
415         if (!process_ok && !pid_ok) {
416                 fprintf(stderr,"Error: You must specify a target process\n");
417                 print_usage (stderr, 1);
418                 exit(1);
419         }
420         if ((exe!=NULL && path!=NULL) || (pid_ok && (exe!=NULL || path!=NULL))) {
421                 fprintf(stderr,"Error: You must specify exactly one target process\n");
422                 print_usage (stderr, 1);
423                 exit(1);
424         }
425         if (!limit_ok) {
426                 fprintf(stderr,"Error: You must specify a cpu limit\n");
427                 print_usage (stderr, 1);
428                 exit(1);
429         }
430         float limit=perclimit/100.0;
431         if (limit<0 || limit >1) {
432                 fprintf(stderr,"Error: limit must be in the range 0-100\n");
433                 print_usage (stderr, 1);
434                 exit(1);
435         }
436         //parameters are all ok!
437         signal(SIGINT,quit);
438         signal(SIGTERM,quit);
439
440         //time quantum in microseconds. it's splitted in a working period and a sleeping one
441         int period=100000;
442         struct timespec twork,tsleep;   //working and sleeping intervals
443         memset(&twork,0,sizeof(struct timespec));
444         memset(&tsleep,0,sizeof(struct timespec));
445
446 wait_for_process:
447
448         //look for the target process..or wait for it
449         if (exe!=NULL)
450                 pid=getpidof(exe);
451         else if (path!=NULL)
452                 pid=getpidof(path);
453         else {
454                 waitforpid(pid);
455         }
456         //process detected...let's play
457
458         //init compute_cpu_usage internal stuff
459         compute_cpu_usage(0,0,NULL);
460         //main loop counter
461         int i=0;
462
463         struct timespec startwork,endwork;
464         long workingtime=0;             //last working time in microseconds
465
466         if (verbose) print_caption();
467
468         float pcpu_avg=0;
469
470         //here we should already have high priority, for time precision
471         while(1) {
472
473                 //estimate how much the controlled process is using the cpu in its working interval
474                 struct cpu_usage cu;
475                 if (compute_cpu_usage(pid,workingtime,&cu)==-1) {
476                         fprintf(stderr,"Process %d dead!\n",pid);
477                         if (lazy) exit(2);
478                         //wait until our process appears
479                         goto wait_for_process;          
480                 }
481
482                 //cpu actual usage of process (range 0-1)
483                 float pcpu=cu.pcpu;
484                 //rate at which we are keeping active the process (range 0-1)
485                 float workingrate=cu.workingrate;
486
487                 //adjust work and sleep time slices
488                 if (pcpu>0) {
489                         twork.tv_nsec=min(period*limit*1000/pcpu*workingrate,period*1000);
490                 }
491                 else if (pcpu==0) {
492                         twork.tv_nsec=period*1000;
493                 }
494                 else if (pcpu==-1) {
495                         //not yet a valid idea of cpu usage
496                         pcpu=limit;
497                         workingrate=limit;
498                         twork.tv_nsec=min(period*limit*1000,period*1000);
499                 }
500                 tsleep.tv_nsec=period*1000-twork.tv_nsec;
501
502                 //update average usage
503                 pcpu_avg=(pcpu_avg*i+pcpu)/(i+1);
504
505                 if (verbose && i%10==0 && i>0) {
506                         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);
507                 }
508
509                 if (limit<1 && limit>0) {
510                         //resume process
511                         if (kill(pid,SIGCONT)!=0) {
512                                 fprintf(stderr,"Process %d dead!\n",pid);
513                                 if (lazy) exit(2);
514                                 //wait until our process appears
515                                 goto wait_for_process;
516                         }
517                 }
518
519                 clock_gettime(CLOCK_REALTIME,&startwork);
520                 nanosleep(&twork,NULL);         //now process is working        
521                 clock_gettime(CLOCK_REALTIME,&endwork);
522                 workingtime=timediff(&endwork,&startwork);
523
524                 if (limit<1) {
525                         //stop process, it has worked enough
526                         if (kill(pid,SIGSTOP)!=0) {
527                                 fprintf(stderr,"Process %d dead!\n",pid);
528                                 if (lazy) exit(2);
529                                 //wait until our process appears
530                                 goto wait_for_process;
531                         }
532                         nanosleep(&tsleep,NULL);        //now process is sleeping
533                 }
534                 i++;
535         }
536
537 }