Load cpulimit/ into cpulimit/branches/upstream/current.
[debian/cpulimit.git] / cpulimit.c
1 /**
2  *
3  * cpulimit - a cpu limiter for Linux
4  *
5  * Copyright (C) 2005-2008, by:  Angelo Marletta <marlonx80@hotmail.com>
6  *
7  * This program is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU General Public License
9  * as published by the Free Software Foundation; either version 2
10  * of the License, or (at your option) any later version.
11  *
12  * This program is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15  * GNU General Public License for more details.
16  *
17  * You should have received a copy of the GNU General Public License
18  * along with this program; if not, write to the Free Software
19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
20  *
21  **********************************************************************
22  *
23  * This is a simple program to limit the cpu usage of a process
24  * If you modify this code, send me a copy please
25  *
26  * Date:    15/2/2008
27  * Version: 1.2 alpha
28  * Get the latest version at: http://cpulimit.sourceforge.net
29  *
30  * Changelog:
31  * - reorganization of the code, splitted in more source files
32  * - control function process_monitor() optimized by eliminating an unnecessary loop
33  * - experimental support for multiple control of children processes and threads
34  *   children detection algorithm seems heavy because of the amount of code,
35  *   but it's designed to be scalable when there are a lot of children processes
36  * - cpu count detection, i.e. if you have 4 cpu, it is possible to limit up to 400%
37  * - in order to avoid deadlock, cpulimit prevents to limit itself
38  * - option --path eliminated, use --exe instead both for absolute path and file name
39  * - deleted almost every setpriority(), just set it once at startup
40  * - minor enhancements and bugfixes
41  *
42  */
43
44
45 #include <getopt.h>
46 #include <stdio.h>
47 #include <fcntl.h>
48 #include <stdlib.h>
49 #include <time.h>
50 #include <sys/time.h>
51 #include <unistd.h>
52 #include <sys/types.h>
53 #include <signal.h>
54 #include <sys/resource.h>
55 #include <string.h>
56 #include <dirent.h>
57 #include <errno.h>
58 #include <string.h>
59
60 #include "process.h"
61 #include "procutils.h"
62 #include "list.h"
63
64 //some useful macro
65 #define MIN(a,b) (a<b?a:b)
66 #define MAX(a,b) (a>b?a:b)
67 #define print_caption() printf("\n%%CPU\twork quantum\tsleep quantum\tactive rate\n")
68
69 //control time slot in microseconds
70 //each slot is splitted in a working slice and a sleeping slice
71 #define TIME_SLOT 100000
72
73 #define MAX_PRIORITY -10
74
75 /* GLOBAL VARIABLES */
76
77 //the "family"
78 struct process_family pf;
79 //pid of cpulimit
80 int cpulimit_pid;
81 //name of this program (maybe cpulimit...)
82 char *program_name;
83
84 /* CONFIGURATION VARIABLES */
85
86 //verbose mode
87 int verbose = 0;
88 //lazy mode (exits if there is no process)
89 int lazy = 0;
90
91 //how many cpu do we have?
92 int get_cpu_count()
93 {
94         FILE *fd;
95         int cpu_count = 0;
96         char line[100];
97         fd = fopen("/proc/stat", "r");
98         if (fd < 0)
99                 return 0; //are we running Linux??
100         while (fgets(line,sizeof(line),fd)!=NULL) {
101                 if (strncmp(line, "cpu", 3) != 0) break;
102                 cpu_count++;
103         }
104         fclose(fd);
105         return cpu_count - 1;
106 }
107
108 //return t1-t2 in microseconds (no overflow checks, so better watch out!)
109 inline unsigned long timediff(const struct timespec *t1,const struct timespec *t2)
110 {
111         return (t1->tv_sec - t2->tv_sec) * 1000000 + (t1->tv_nsec/1000 - t2->tv_nsec/1000);
112 }
113
114 //returns t1-t2 in microseconds
115 inline unsigned long long tv_diff(struct timeval *t1, struct timeval *t2)
116 {
117         return ((unsigned long long)(t1->tv_sec - t2->tv_sec)) * 1000000ULL + t1->tv_usec - t2->tv_usec;
118 }
119
120 //SIGINT and SIGTERM signal handler
121 void quit(int sig)
122 {
123         //let all the processes continue if stopped
124         struct list_node *node = NULL;
125         for (node=pf.members.first; node!= NULL; node=node->next) {
126                 struct process *p = (struct process*)(node->data);
127                 process_close(p->history);
128                 kill(p->pid, SIGCONT);
129         }
130         //free all the memory
131         cleanup_process_family(&pf);
132         exit(0);
133 }
134
135 void print_usage(FILE *stream, int exit_code)
136 {
137         fprintf(stream, "Usage: %s TARGET [OPTIONS...]\n",program_name);
138         fprintf(stream, "   TARGET must be exactly one of these:\n");
139         fprintf(stream, "      -p, --pid=N        pid of the process (implies -z)\n");
140         fprintf(stream, "      -e, --exe=FILE     name of the executable program file or absolute path name\n");
141         fprintf(stream, "   OPTIONS\n");
142         fprintf(stream, "      -l, --limit=N      percentage of cpu allowed from 0 to 100 (required)\n");
143         fprintf(stream, "      -v, --verbose      show control statistics\n");
144         fprintf(stream, "      -z, --lazy         exit if there is no suitable target process, or if it dies\n");
145         fprintf(stream, "      -h, --help         display this help and exit\n");
146         exit(exit_code);
147 }
148
149 void limit_process(int pid, double limit)
150 {
151         //slice of the slot in which the process is allowed to run
152         struct timespec twork;
153         //slice of the slot in which the process is stopped
154         struct timespec tsleep;
155         //when the last twork has started
156         struct timespec startwork;
157         //when the last twork has finished
158         struct timespec endwork;
159         //initialization
160         memset(&twork, 0, sizeof(struct timespec));
161         memset(&tsleep, 0, sizeof(struct timespec));
162         memset(&startwork, 0, sizeof(struct timespec));
163         memset(&endwork, 0, sizeof(struct timespec));   
164         //last working time in microseconds
165         unsigned long workingtime = 0;
166         int i = 0;
167
168         //build the family
169         create_process_family(&pf, pid);
170         struct list_node *node;
171         
172         if (verbose) printf("Members in the family owned by %d: %d\n", pf.father, pf.members.count);
173
174         //rate at which we are keeping active the processes (range 0-1)
175         //1 means that the process are using all the twork slice
176         double workingrate = -1;
177
178         while(1) {
179
180                 if (i%200==0 && verbose) print_caption();
181
182                 if (i%10==0) {
183                         //update the process family (checks only for new members)
184                         int new_children = update_process_family(&pf);
185                         if (new_children) {
186                                 printf("%d new children processes detected (", new_children);
187                                 int j;
188                                 node = pf.members.last;
189                                 for (j=0; j<new_children; j++) {
190                                         printf("%d", ((struct process*)(node->data))->pid);
191                                         if (j<new_children-1) printf(" ");
192                                         node = node->previous;
193                                 }
194                                 printf(")\n");
195                         }
196                 }
197
198                 //total cpu actual usage (range 0-1)
199                 //1 means that the processes are using 100% cpu
200                 double pcpu = -1;
201                 //number of processes in the family
202                 int pcount = 0;
203                 
204                 //estimate how much the controlled processes are using the cpu in the working interval
205                 for (node=pf.members.first; node!=NULL; node=node->next) {
206                         struct process *proc = (struct process*)(node->data);
207                         if (process_monitor(proc->history)==-1) {
208                                 //process is dead, remove it from family
209                                 fprintf(stderr,"Process %d dead!\n", proc->pid);
210                                 remove_process_from_family(&pf, proc->pid);
211                                 continue;
212                         }
213 //printf("pid %d limit %f pcpu %f wrate %f\n", proc->pid, limit, proc->history->usage.pcpu, proc->history->usage.workingrate);
214                         if (proc->history->cpu_usage<0) {
215                                 continue;
216                         }
217                         if (pcpu<0) pcpu = 0;
218                         pcpu += proc->history->cpu_usage;
219                         pcount++;
220                 }
221
222                 //adjust work and sleep time slices
223                 if (pcpu < 0) {
224                         //it's the 1st cycle, initialize workingrate
225                         pcpu = limit;
226                         workingrate = limit;
227                         twork.tv_nsec = TIME_SLOT*limit*1000;
228                 }
229                 else {
230                         //adjust workingrate
231                         workingrate = MIN(workingrate / pcpu * limit, 1);
232                         twork.tv_nsec = TIME_SLOT*1000*workingrate;
233                 }
234                 tsleep.tv_nsec = TIME_SLOT*1000-twork.tv_nsec;
235
236 //printf("%lf %lf\n", workingrate, pcpu);
237
238                 if (verbose && i%10==0 && i>0) {
239                         printf("%0.2lf%%\t%6ld us\t%6ld us\t%0.2lf%%\n",pcpu*100,twork.tv_nsec/1000,tsleep.tv_nsec/1000,workingrate*100);
240                 }
241
242                 //resume processes
243                 for (node=pf.members.first; node!=NULL; node=node->next) {
244                         struct process *proc = (struct process*)(node->data);
245                         if (kill(proc->pid,SIGCONT)!=0) {
246                                 //process is dead, remove it from family
247                                 fprintf(stderr,"Process %d dead!\n", proc->pid);
248                                 remove_process_from_family(&pf, proc->pid);
249                         }
250                 }
251
252                 //now processes are free to run (same working slice for all)
253                 clock_gettime(CLOCK_REALTIME,&startwork);
254                 nanosleep(&twork,NULL);
255                 clock_gettime(CLOCK_REALTIME,&endwork);
256                 workingtime = timediff(&endwork,&startwork);
257
258                 if (tsleep.tv_nsec>0) {
259                         //stop only if tsleep>0, instead it's useless
260                         for (node=pf.members.first; node!=NULL; node=node->next) {
261                                 struct process *proc = (struct process*)(node->data);
262                                 if (kill(proc->pid,SIGSTOP)!=0) {
263                                         //process is dead, remove it from family
264                                         fprintf(stderr,"Process %d dead!\n", proc->pid);
265                                         remove_process_from_family(&pf, proc->pid);
266                                 }
267                         }
268                         //now the processes are sleeping
269                         nanosleep(&tsleep,NULL);
270                 }
271                 i++;
272         }
273         cleanup_process_family(&pf);
274 }
275
276 int main(int argc, char **argv) {
277         //get program name
278         char *p=(char*)memrchr(argv[0],(unsigned int)'/',strlen(argv[0]));
279         program_name = p==NULL?argv[0]:(p+1);
280         cpulimit_pid = getpid();
281
282         //argument variables
283         const char *exe = NULL;
284         int perclimit = 0;
285         int pid_ok = 0;
286         int process_ok = 0;
287         int limit_ok = 0;
288         int pid = 0;
289
290         //parse arguments
291         int next_option;
292     int option_index = 0;
293         //A string listing valid short options letters
294         const char* short_options = "p:e:l:vzh";
295         //An array describing valid long options
296         const struct option long_options[] = {
297                 { "pid",        required_argument, NULL,     'p' },
298                 { "exe",        required_argument, NULL,     'e' },
299                 { "limit",      required_argument, NULL,     'l' },
300                 { "verbose",    no_argument,       &verbose, 'v' },
301                 { "lazy",       no_argument,       &lazy,    'z' },
302                 { "help",       no_argument,       NULL,     'h' },
303                 { 0,            0,                 0,         0  }
304         };
305
306         do {
307                 next_option = getopt_long(argc, argv, short_options,long_options, &option_index);
308                 switch(next_option) {
309                         case 'p':
310                                 pid = atoi(optarg);
311                                 //todo: verify pid is valid
312                                 pid_ok = 1;
313                                 process_ok = 1;
314                                 break;
315                         case 'e':
316                                 exe = optarg;
317                                 process_ok = 1;
318                                 break;
319                         case 'l':
320                                 perclimit = atoi(optarg);
321                                 limit_ok = 1;
322                                 break;
323                         case 'v':
324                                 verbose = 1;
325                                 break;
326                         case 'z':
327                                 lazy = 1;
328                                 break;
329                         case 'h':
330                                 print_usage(stdout, 1);
331                                 break;
332                         case '?':
333                                 print_usage(stderr, 1);
334                                 break;
335                         case -1:
336                                 break;
337                         default:
338                                 abort();
339                 }
340         } while(next_option != -1);
341
342         if (pid!=0) {
343                 lazy = 1;
344         }
345         
346         if (pid_ok && (pid<=1 || pid>=65536)) {
347                 fprintf(stderr,"Error: Invalid value for argument PID\n");
348                 print_usage(stderr, 1);
349                 exit(1);
350         }
351         
352         if (!process_ok) {
353                 fprintf(stderr,"Error: You must specify a target process, either by name or by PID\n");
354                 print_usage(stderr, 1);
355                 exit(1);
356         }
357         if (pid_ok && exe!=NULL) {
358                 fprintf(stderr, "Error: You must specify exactly one process, either by name or by PID\n");
359                 print_usage(stderr, 1);
360                 exit(1);
361         }
362         if (!limit_ok) {
363                 fprintf(stderr,"Error: You must specify a cpu limit percentage\n");
364                 print_usage(stderr, 1);
365                 exit(1);
366         }
367         double limit = perclimit/100.0;
368         int cpu_count = get_cpu_count();
369         printf("%d cpu detected\n", cpu_count);
370         if (limit<0 || limit >cpu_count) {
371                 fprintf(stderr,"Error: limit must be in the range 0-%d00\n", cpu_count);
372                 print_usage(stderr, 1);
373                 exit(1);
374         }
375         //parameters are all ok!
376         signal(SIGINT, quit);
377         signal(SIGTERM, quit);
378
379         //try to renice with the best value
380         int old_priority = getpriority(PRIO_PROCESS, 0);
381         int priority = old_priority;
382         while (setpriority(PRIO_PROCESS, 0, priority-1) == 0 && priority>MAX_PRIORITY) {
383                 priority--;     
384         }
385         if (priority != old_priority) {
386                 printf("Priority changed to %d\n", priority);
387         }
388         else {
389                 printf("Warning: Cannot change priority. Run as root for best results.\n");
390         }
391         
392         while(1) {
393                 //look for the target process..or wait for it
394                 int ret = 0;
395                 if (pid_ok) {
396                         //search by pid
397                         ret = look_for_process_by_pid(pid);
398                         if (ret == 0) {
399                                 printf("No process found\n");
400                         }
401                         else if (ret < 0) {
402                                 printf("Process found but you aren't allowed to control it\n");
403                         }
404                 }
405                 else {
406                         //search by file or path name
407                         ret = look_for_process_by_name(exe);
408                         if (ret == 0) {
409                                 printf("No process found\n");
410                         }
411                         else if (ret < 0) {
412                                 printf("Process found but you aren't allowed to control it\n");
413                         }
414                         else {
415                                 pid = ret;
416                         }
417                 }
418                 if (ret > 0) {
419                         if (ret == cpulimit_pid) {
420                                 printf("Process %d is cpulimit itself! Aborting to avoid deadlock\n", ret);
421                                 exit(1);
422                         }
423                         printf("Process %d found\n", pid);
424                         //control
425                         limit_process(pid, limit);
426                 }
427                 if (lazy) {
428                         printf("Giving up...\n");
429                         break;
430                 }
431                 sleep(2);
432         }
433         
434         return 0;
435 }
436