New upstream version 2.4
[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  * Date: Jan 29, 2013
19  * Version 1.2 and newer
20  *
21  * Modifications and updates by: Hasnain Lakhani
22  * Date: Mar 26, 2014
23  * Version 2.1
24  */
25
26
27 #include <getopt.h>
28 #include <stdio.h>
29 #include <stdlib.h>
30 #include <time.h>
31 #include <sys/time.h>
32 #include <sys/types.h>
33 #include <sys/wait.h>
34 #include <unistd.h>
35 #include <sys/types.h>
36 #include <signal.h>
37 #include <sys/resource.h>
38 #include <string.h>
39 #include <dirent.h>
40 #include <errno.h>
41 #include <string.h>
42 #include <limits.h>    // for compatibility
43
44 #ifdef __APPLE__
45 #include <mach/clock.h>
46 #include <mach/mach.h>
47 #endif
48
49
50 #ifdef FREEBSD
51 #include <fcntl.h>
52 #include <kvm.h>
53 #include <paths.h>
54 #include <sys/param.h>
55 #include <sys/sysctl.h>
56 #include <sys/user.h>
57 #endif
58
59 #ifdef LINUX
60 #include <dirent.h>
61 #define PROC_FILENAME 64
62 #define LINE_LENGTH 256
63 #endif
64
65 //kernel time resolution (inverse of one jiffy interval) in Hertz
66 //i don't know how to detect it, then define to the default (not very clean!)
67 #define HZ 100
68
69 //some useful macro
70 #define min(a,b) (a<b?a:b)
71 #define max(a,b) (a>b?a:b)
72
73 // For platforms without PATH_MAX
74 #ifndef PATH_MAX
75 #define PATH_MAX 4096
76 #endif
77
78 #define BEST_PRIORITY -10
79
80 #ifndef TRUE
81 #define TRUE 1
82 #endif
83 #ifndef FALSE
84 #define FALSE 0
85 #endif
86
87 #ifndef VERSION
88 #define VERSION 2.4
89 #endif
90
91 //pid of the controlled process
92 pid_t pid = 0;
93 pid_t my_pid;     // this process's PID
94
95 //executable file name
96 char *program_name;
97 //verbose mode
98 int verbose = FALSE;
99 //lazy mode
100 int lazy = FALSE;
101 // is higher priority nice possible?
102 int nice_lim;
103
104 // number of CPUs we detected
105 int NCPU;
106
107 // quiet mode
108 int quiet = FALSE;
109
110 // What signal should we send to the watched process
111 // when cpulimit exits?
112 int send_signal = SIGCONT;
113
114 //reverse byte search
115 // void *memrchr(const void *s, int c, size_t n);
116
117 #define MAX_SIGNAL 7
118 const char *SIGNAL_NAME[MAX_SIGNAL] = { "SIGHUP", "SIGINT", "SIGQUIT", 
119                                  "SIGKILL", "SIGTERM", "SIGSTOP", "SIGCONT" };
120 const int SIGNAL_VALUE[MAX_SIGNAL] = { SIGHUP, SIGINT, SIGQUIT,
121                                   SIGKILL, SIGTERM, SIGSTOP, SIGCONT };
122
123
124 //return ta-tb in microseconds (no overflow checks!)
125 inline long timediff(const struct timespec *ta,const struct timespec *tb) {
126     unsigned long us = (ta->tv_sec-tb->tv_sec)*1000000 + (ta->tv_nsec/1000 - tb->tv_nsec/1000);
127     return us;
128 }
129
130
131 #ifdef LINUX
132 #include <pthread.h>
133
134 typedef struct
135 {
136    pid_t child;   // the child of our target process
137    pid_t monitor; // the LimitCPU fork monitoring the child
138    void *next; 
139 } CHILD;
140
141 int quitting = FALSE;           // Have we receive a quit signal
142 int monitor_children = FALSE;   // are we monitoring children of target
143
144 // Data passed to the monitor thread
145 typedef struct
146 {
147    int limit;           // per cent limit to place on a process
148    char *this_program;  // copy of argv[0]
149 } PROGRAM_DATA;
150
151 #endif
152
153
154
155 int Check_Us(pid_t target_pid)
156 {
157    pid_t this_pid;
158
159    this_pid = getpid();
160    if (this_pid == target_pid)
161    {
162       fprintf(stderr, "We cannot throttle ourselves.\n");
163       exit(7);
164    }
165    return TRUE;
166 }
167
168
169 int waitforpid(int pid) {
170         //switch to low priority
171         // if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
172         /*
173         if ( (nice_lim < INT_MAX) && 
174              (setpriority(PRIO_PROCESS, my_pid, 19) != 0) ) {
175                 printf("Warning: cannot renice\n");
176         }
177         */
178         int i=0;
179
180         while(1) {
181
182                 DIR *dip;
183                 struct dirent *dit;
184
185                 //open a directory stream to /proc directory
186                 if ((dip = opendir("/proc")) == NULL) {
187                         perror("opendir");
188                         return -1;
189                 }
190
191                 //read in from /proc and seek for process dirs
192                 while ((dit = readdir(dip)) != NULL) {
193                         //get pid
194                         if (pid==atoi(dit->d_name)) {
195                                 //pid detected
196                                 Check_Us(pid);
197                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
198                                         //process is ok!
199                                         if (closedir(dip) == -1) {
200                                            perror("closedir");
201                                            return -1;
202                                         }
203                                         goto done;
204                                 }
205                                 else {
206                                         fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
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                                 fprintf(stderr, "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     if (!quiet)
234         printf("Process %d detected\n",pid);
235         //now set high priority, if possible
236         // if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
237         /*
238         if ( (nice_lim < INT_MAX) &&
239              (setpriority(PRIO_PROCESS, my_pid, nice_lim) != 0) ) {
240                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
241         }
242         */
243         return 0;
244
245 }
246
247 //this function periodically scans process list and looks for executable path names
248 //it should be executed in a low priority context, since precise timing does not matter
249 //if a process is found then its pid is returned
250 //process: the name of the wanted process, can be an absolute path name to the executable file
251 //         or simply its name
252 //return: pid of the found process
253 int getpidof(const char *process) {
254
255         //set low priority
256         // if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
257         /*
258         if ( (nice_lim < INT_MAX) &&
259              (setpriority(PRIO_PROCESS, my_pid, 19) != 0) ) {
260                 printf("Warning: cannot renice\n");
261         }
262         */
263         char exelink[20];
264         char exepath[PATH_MAX+1];
265         int pid=0;
266         int i=0;
267
268         while(1) {
269
270                 DIR *dip;
271                 struct dirent *dit;
272
273                 //open a directory stream to /proc directory
274                 if ((dip = opendir("/proc")) == NULL) {
275                         perror("opendir");
276                         return -1;
277                 }
278
279                 //read in from /proc and seek for process dirs
280                 while ((dit = readdir(dip)) != NULL) {
281                         //get pid
282                         pid=atoi(dit->d_name);
283                         if (pid>0) {
284                                 sprintf(exelink,"/proc/%d/exe",pid);
285                                 int size=readlink(exelink,exepath,sizeof(exepath));
286                                 if (size>0) {
287                                         int found=0;
288                                         if (process[0]=='/' && strncmp(exepath,process,size)==0 && size==strlen(process)) {
289                                                 //process starts with / then it's an absolute path
290                                                 found=1;
291                                         }
292                                         else {
293                                                 //process is the name of the executable file
294                                                 if (strncmp(exepath+size-strlen(process),process,strlen(process))==0) {
295                                                         found=1;
296                                                 }
297                                         }
298                                         if (found==1) {
299                                         Check_Us(pid);
300                                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
301                                                         //process is ok!
302                                                         if (closedir(dip) == -1) {
303                                                           perror("closedir");
304                                                           return -1;
305                                                         }
306                                                         goto done;
307                                                 }
308                                                 else {
309                                                         fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
310                                                 }
311                                         }
312                                 }
313                         }
314                 }
315
316                 //close the dir stream and check for errors
317                 if (closedir(dip) == -1) {
318                         perror("closedir");
319                         return -1;
320                 }
321
322                 //no suitable target found
323                 if (i++==0) {
324                         if (lazy) {
325                                 fprintf(stderr,"No process found\n");
326                                 exit(2);
327                         }
328                         else {
329                                 fprintf(stderr, "Warning: no target process found. Waiting for it...\n");
330                         }
331                 }
332
333                 //sleep for a while
334                 sleep(2);
335         }
336
337 done:
338     if (!quiet)
339         printf("Process %d detected\n",pid);
340         //now set high priority, if possible
341         // if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
342         /*
343         if ( (nice_lim < INT_MAX) &&
344              (setpriority(PRIO_PROCESS, my_pid, nice_lim) != 0) ) {
345                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
346         }
347         */
348         return pid;
349
350 }
351
352 //SIGINT and SIGTERM signal handler
353 void quit(int sig) {
354         //let the process continue if we are stopped
355         kill(pid, send_signal);
356         #ifdef LINUX
357         if (monitor_children)
358         {
359             quitting = TRUE;
360             printf("Asking children to quit...\n");
361             sleep(2);                // wait for thread clean-up
362         }
363         #endif
364         printf("Exiting...\n");
365         exit(0);
366 }
367
368 // Handle a child process quitting
369 void Child_Done(int sig)
370 {
371    pid_t caught_child;
372    caught_child = waitpid(-1, NULL, WNOHANG);
373    if (verbose)
374    {
375       printf("Caught child process: %d\n", (int) caught_child);
376       printf("%d\n", errno);
377    }
378    // If this was the one process we were watching, we can quit now.
379    if (caught_child == pid)
380    {
381       printf("Child process is finished, exiting...\n");
382       exit(0);
383    }
384 }
385
386
387 #ifdef FREEBSD
388 int getjiffies(int pid)
389 {
390    kvm_t *my_kernel = NULL;
391    struct kinfo_proc *process_data = NULL;
392    int processes;
393    int my_jiffies = -1;
394
395    my_kernel = kvm_open(0, 0, 0, O_RDONLY, "kvm_open");
396    if (! my_kernel)
397    {
398       fprintf(stderr, "Error opening kernel vm. You should be running as root.\n");
399       return -1;
400    }
401
402    process_data = kvm_getprocs(my_kernel, KERN_PROC_PID, pid, &processes);
403    if ( (process_data) && (processes >= 1) )
404        my_jiffies = process_data->ki_runtime;
405    
406    kvm_close(my_kernel);
407    if (my_jiffies >= 0)
408      my_jiffies /= 1000;
409    return my_jiffies;
410 }
411
412 #endif
413
414 #ifdef LINUX
415 //get jiffies count from /proc filesystem
416 int getjiffies(int pid) {
417         static char stat[20];
418         static char buffer[1024];
419         char *p;
420         sprintf(stat,"/proc/%d/stat",pid);
421         FILE *f=fopen(stat,"r");
422         if (f==NULL) return -1;
423         p = fgets(buffer,sizeof(buffer),f);
424         fclose(f);
425         // char *p=buffer;
426         if (p)
427         {
428           p=memchr(p+1,')',sizeof(buffer)-(p-buffer));
429           int sp=12;
430           while (sp--)
431                 p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
432           //user mode jiffies
433           int utime=atoi(p+1);
434           p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
435           //kernel mode jiffies
436           int ktime=atoi(p+1);
437           return utime+ktime;
438         }
439         // could not read info
440         return -1;
441 }
442 #endif
443
444
445 //process instant photo
446 struct process_screenshot {
447         struct timespec when;   //timestamp
448         int jiffies;    //jiffies count of the process
449         int cputime;    //microseconds of work from previous screenshot to current
450 };
451
452 //extracted process statistics
453 struct cpu_usage {
454         float pcpu;
455         float workingrate;
456 };
457
458 //this function is an autonomous dynamic system
459 //it works with static variables (state variables of the system), that keep memory of recent past
460 //its aim is to estimate the cpu usage of the process
461 //to work properly it should be called in a fixed periodic way
462 //perhaps i will put it in a separate thread...
463 int compute_cpu_usage(int pid,int last_working_quantum,struct cpu_usage *pusage) {
464         #define MEM_ORDER 10
465         //circular buffer containing last MEM_ORDER process screenshots
466         static struct process_screenshot ps[MEM_ORDER];
467         //the last screenshot recorded in the buffer
468         static int front=-1;
469         //the oldest screenshot recorded in the buffer
470         static int tail=0;
471
472         if (pusage==NULL) {
473                 //reinit static variables
474                 front=-1;
475                 tail=0;
476                 return 0;
477         }
478
479         //let's advance front index and save the screenshot
480         front=(front+1)%MEM_ORDER;
481         int j=getjiffies(pid);
482         if (j>=0) ps[front].jiffies=j;
483         else return -1; //error: pid does not exist
484
485         #ifdef __APPLE__
486         // OS X does not have clock_gettime, use clock_get_time
487         clock_serv_t cclock;
488         mach_timespec_t mts;
489         host_get_clock_service(mach_host_self(), CALENDAR_CLOCK, &cclock);
490         clock_get_time(cclock, &mts);
491         mach_port_deallocate(mach_task_self(), cclock);
492         ps[front].when.tv_sec = mts.tv_sec;
493         ps[front].when.tv_nsec = mts.tv_nsec;
494
495         #else
496         // Linux and BSD can use real time
497         clock_gettime(CLOCK_REALTIME,&(ps[front].when));
498         ps[front].cputime=last_working_quantum;
499         #endif
500         //buffer actual size is: (front-tail+MEM_ORDER)%MEM_ORDER+1
501         int size=(front-tail+MEM_ORDER)%MEM_ORDER+1;
502
503         if (size==1) {
504                 //not enough samples taken (it's the first one!), return -1
505                 pusage->pcpu=-1;
506                 pusage->workingrate=1;
507                 return 0;
508         }
509         else {
510                 //now we can calculate cpu usage, interval dt and dtwork are expressed in microseconds
511                 long dt=timediff(&(ps[front].when),&(ps[tail].when));
512                 long dtwork=0;
513                 int i=(tail+1)%MEM_ORDER;
514                 int max=(front+1)%MEM_ORDER;
515                 do {
516                         dtwork+=ps[i].cputime;
517                         i=(i+1)%MEM_ORDER;
518                 } while (i!=max);
519                 int used=ps[front].jiffies-ps[tail].jiffies;
520                 float usage=(used*1000000.0/HZ)/dtwork;
521                 pusage->workingrate=1.0*dtwork/dt;
522                 pusage->pcpu=usage*pusage->workingrate;
523                 if (size==MEM_ORDER)
524                         tail=(tail+1)%MEM_ORDER;
525                 return 0;
526         }
527         #undef MEM_ORDER
528 }
529
530 void print_caption() {
531         printf("\n%%CPU\twork quantum\tsleep quantum\tactive rate\n");
532 }
533
534
535 void increase_priority()
536 {
537         //find the best available nice value
538         int old_priority = getpriority(PRIO_PROCESS, 0);
539         int priority = old_priority;
540         while ( (setpriority(PRIO_PROCESS, 0, priority-1) == 0) &&
541                 (priority > BEST_PRIORITY) )
542         {
543                 priority--;
544         }
545         if (priority != old_priority) {
546                 if (verbose) printf("Priority changed to %d\n", priority);
547         }
548         else {
549                 if (verbose) printf("Warning: Cannot change priority. Run as root or renice for best results.\n");
550         }
551
552
553 }
554
555
556
557 #ifdef LINUX
558 // This following functions are for detecting and limiting child
559 // processes on Linux.
560
561 // This function adds a new child process to our list of child processes.
562 CHILD *Add_Child(CHILD *all_children, pid_t new_pid)
563 {
564     CHILD *new_child = (CHILD *) calloc(sizeof(CHILD), 1);
565     CHILD *current;
566
567     if (! new_child)
568        return all_children;
569
570     new_child->next = NULL;
571     new_child->child = new_pid;
572     new_child->monitor = 0;
573
574     if (all_children)
575     {
576         current = all_children;
577         while (current->next)
578            current = current->next;
579         current->next = new_child;
580         return all_children;
581     }
582     else   // this is the first node
583        return new_child;
584 }
585
586
587
588 // This function removes a child PID node.
589 CHILD *Remove_Child(CHILD *all_children, pid_t old_pid)
590 {
591     CHILD *current, *previous = NULL;
592     int found = FALSE;
593
594     current = all_children;
595     while ( (! found) && (current) )
596     {
597         if (current->child == old_pid)
598         {
599            if (previous)
600                previous->next = current->next;
601            else
602              all_children = current->next;
603            
604            free(current);
605            found = TRUE;
606         }
607         else
608         {
609             previous = current;
610             current = current->next;
611         }
612      }
613      return all_children;
614 }
615
616
617 // This function cleans up all remaining child nodes.
618 void Clean_Up_Children(CHILD *all_children)
619 {
620    CHILD *current, *next;
621
622    current = all_children;
623    while (current)
624    {
625       next = current->next;
626       free(current);
627       current = next;
628    }
629 }
630
631
632 // This function searches the linked list for a matching PID.
633 // It returns NULL if no match is found and a pointer to the
634 // node is a match is located.
635 CHILD *Find_Child(CHILD *children, pid_t target)
636 {
637    CHILD *current;
638    int found = FALSE;
639
640    current = children;
641    while ( (!found) && (current) )
642    {
643       if (current->child == target)
644         found = TRUE;
645       else
646         current = current->next;
647    }
648    return current;
649 }
650
651
652
653 // This function returns a list of process IDs of children
654 // of the given process (PID). It does this by searching the /proc
655 // file system and looking in the /proc/pid/status file for the PPid field.
656 // A linked list of child PIDs is returned on success or NULL on failure or
657 // if no child PIDs are found.
658 CHILD *Find_Child_PIDs(CHILD *all_children, pid_t parent_pid)
659 {
660     int found = FALSE;
661     DIR *proc;
662     struct dirent *proc_entry;
663     char filename[PROC_FILENAME];
664     FILE *status_file;
665     char *reading_file;
666     char line[256];
667     pid_t new_ppid;
668     int current_pid;
669
670     proc = opendir("/proc");
671     if (! proc)
672        return all_children;
673
674     proc_entry = readdir(proc);
675     while (proc_entry)
676     {
677         snprintf(filename, PROC_FILENAME, "/proc/%s/status", proc_entry->d_name);
678         status_file = fopen(filename, "r");
679         if (status_file)
680         {
681            found = FALSE;
682            reading_file = fgets(line, LINE_LENGTH, status_file);
683            while ( (! found) && (reading_file) )
684            {
685              if (! strncmp(line, "PPid:", 5) )
686              {
687                  sscanf(&(line[6]), "%d", &new_ppid);
688                  if (new_ppid == parent_pid)
689                  {
690                      sscanf(proc_entry->d_name, "%d", &current_pid);
691                      if (! Find_Child(all_children, current_pid) )
692                         all_children = Add_Child(all_children, current_pid);
693                  }
694                  found = TRUE;
695              }
696              else
697                 reading_file = fgets(line, LINE_LENGTH, status_file);
698            }   // done reading status file
699           
700            fclose(status_file);
701         }
702         proc_entry = readdir(proc);
703     }   // done reading proc file system
704     closedir(proc);
705     return all_children;
706 }
707
708
709 // This function (which should probably be called as a thread) monitors the
710 // system for child processes of the current target process. When a new
711 // child of the target is located, it is added to the CHILD list.
712 // New children result in a new fork of this program being spawned to
713 // monitor the child process and its children.
714
715 void *Monitor_Children(void *all_data)
716 {
717    CHILD *all_children = NULL;
718    CHILD *current;
719    PROGRAM_DATA *program_data = (PROGRAM_DATA *) all_data;
720
721    while (! quitting )
722    {
723       // Check for new child processes
724       all_children = Find_Child_PIDs(all_children, pid);
725
726       // Find any children without monitors and create a monitoring process
727       // Clean out old processes while we are looking
728       current = all_children;
729       while (current)
730       {
731          // First see if the child process is still running. If not,
732          // we can remote its node.
733          if (current->child)
734          {
735              char filename[PROC_FILENAME];
736              DIR *child_directory;
737              snprintf(filename, PROC_FILENAME, "/proc/%d", current->child);
738              child_directory = opendir(filename);
739              if (child_directory)
740                 closedir(child_directory);
741              else
742              {
743                 if (verbose)
744                   printf("Child process %d done, cleaning up.\n",
745                           (int) current->child);
746                 all_children = Remove_Child(all_children, current->child);
747              }
748          }  // end of clean up children processes no longer running
749
750          // The child process is still running, but it might not have
751          // a monitor. Create a new monitoring process.
752          if ( (current->child) && (! current->monitor) )
753          {
754               pid_t returned_pid;
755
756               if (verbose)
757                  printf("Creating monitoring process for %d\n",
758                         (int) current->child);
759
760               returned_pid = fork();
761               if (returned_pid > 0)
762               {
763                  // parent
764                  current->monitor = returned_pid;
765               }
766               else if (returned_pid == 0)
767               {
768                  // child
769                  char limit_amount[16];
770                  char process_identifier[16];
771                  snprintf(limit_amount, 16, "%d", (int) program_data->limit);
772                  snprintf(process_identifier, 16, "%d", current->child);
773                  if (verbose)
774                     printf("Starting monitor with: %s -l %s -p %s -z -m\n",
775                             program_data->this_program, limit_amount,
776                             process_identifier);
777                  execl(program_data->this_program, program_data->this_program,
778                        "-l", limit_amount, "-p", process_identifier, 
779                        "-z", "-m", (char *) NULL);
780               }
781               
782          }    // end of creating a new monitor
783          if (verbose)
784          {
785              printf("Watching child: %d with %d\n", 
786                    (int) current->child, (int) current->monitor);
787          }
788          current = current->next;
789       }
790       sleep(1);
791    }  // end LimitCPU is still running
792    pthread_exit(NULL);
793 }
794
795 #endif       // end of monitoring children processes on Linux
796
797
798
799 void print_usage(FILE *stream,int exit_code) {
800         fprintf(stream, "CPUlimit version %1.1f\n", VERSION);
801         fprintf(stream, "Usage: %s TARGET [OPTIONS...] [-- PROGRAM]\n",program_name);
802         fprintf(stream, "   TARGET must be exactly one of these:\n");
803         fprintf(stream, "      -p, --pid=N        pid of the process\n");
804         fprintf(stream, "      -e, --exe=FILE     name of the executable program file\n");
805         fprintf(stream, "                         The -e option only works when\n");
806         fprintf(stream, "                         cpulimit is run with admin rights.\n");
807         fprintf(stream, "      -P, --path=PATH    absolute path name of the\n");
808         fprintf(stream, "                         executable program file\n");
809         fprintf(stream, "   OPTIONS\n");
810         fprintf(stream, "      -b  --background   run in background\n");
811         fprintf(stream, "      -f  --foreground   launch target process in foreground and wait for it to exit\n");
812         fprintf(stream, "      -c  --cpu=N        override the detection of CPUs on the machine.\n");
813         fprintf(stream, "      -l, --limit=N      percentage of cpu allowed from 1 up.\n");
814         fprintf(stream, "                         Usually 1 - %d00, but can be higher\n", NCPU);
815         fprintf(stream, "                         on multi-core CPUs (mandatory)\n");
816         #ifdef LINUX
817         fprintf(stream, "      -m, --monitor-forks  Watch children/forks of the target process\n");
818         #endif
819         fprintf(stream, "      -q, --quiet        run in quiet mode (only print errors).\n");
820         fprintf(stream, "      -k, --kill         kill processes going over their limit\n");
821         fprintf(stream, "                         instead of just throttling them.\n");
822         fprintf(stream, "      -r, --restore      Restore processes after they have\n");
823         fprintf(stream, "                         been killed. Works with the -k flag.\n");
824         fprintf(stream, "      -s, --signal=SIG   Send this signal to the watched process when cpulimit exits.\n");
825         fprintf(stream, "                         Signal should be specificed as a number or \n");
826         fprintf(stream, "                         SIGTERM, SIGCONT, SIGSTOP, etc. SIGCONT is the default.\n");
827         fprintf(stream, "      -v, --verbose      show control statistics\n");
828         fprintf(stream, "      -z, --lazy         exit if there is no suitable target process,\n");
829         fprintf(stream, "                         or if it dies\n");
830         fprintf(stream, "          --             This is the final CPUlimit option. All following\n");
831         fprintf(stream, "                         options are for another program we will launch.\n");
832         fprintf(stream, "      -h, --help         display this help and exit\n");
833         exit(exit_code);
834 }
835
836
837
838 // Get the number of CPU cores on this machine.
839 int get_ncpu()
840 {
841         int ncpu = 1;
842 #ifdef _SC_NPROCESSORS_ONLN
843         ncpu = sysconf(_SC_NPROCESSORS_ONLN);
844 #endif
845         return ncpu;
846 }
847
848
849 // This function attempts to figure out what signal we should send
850 // target processes based on a command line paramter. First we check
851 // for text such as SIGINT, SIGCONT, SIGSTOP, etc. If no match is found
852 // then we assume the value given is a number and use that.
853 int Translate_Signal(char *my_signal)
854 {
855     int signal_value;
856     int index = 0, found = FALSE;
857     // first check to see if we were passed a string
858     while ( (index < MAX_SIGNAL) && (! found) )
859     {
860         if (! strcmp(my_signal, SIGNAL_NAME[index]) )
861         {
862             found = TRUE;
863             signal_value = SIGNAL_VALUE[index];
864         }
865         else
866            index++;
867     }
868
869     // no value found, try a number
870     if (! found)
871        signal_value = atoi(my_signal);
872
873     return signal_value;
874 }
875
876
877
878 int main(int argc, char **argv) {
879
880         //get program name
881         // char *p=(char*)memrchr(argv[0],(unsigned int)'/',strlen(argv[0]));
882         // program_name = p==NULL?argv[0]:(p+1);
883         program_name = argv[0];
884         int run_in_background = FALSE;
885         //parse arguments
886         int next_option;
887         /* A string listing valid short options letters. */
888         #ifdef LINUX
889         const char* short_options="p:e:P:l:c:s:bfqkmrvzh";
890         PROGRAM_DATA program_data;
891         #else
892         const char* short_options="p:e:P:l:c:s:bfqkrvzh";
893         #endif
894         /* An array describing valid long options. */
895         const struct option long_options[] = {
896                 { "pid", required_argument, NULL, 'p' },
897                 { "exe", required_argument, NULL, 'e' },
898                 { "path", required_argument, NULL, 'P' },
899                 { "limit", required_argument, NULL, 'l' },
900                 { "background", no_argument, NULL, 'b' },
901                 { "foreground", no_argument, NULL, 'f' },
902                 { "quiet", no_argument, NULL, 'q' },
903                 { "verbose", no_argument, NULL, 'v' },
904                 { "lazy", no_argument, NULL, 'z' },
905                 { "help", no_argument, NULL, 'h' },
906                 { "cpu", required_argument, NULL, 'c'},
907                 { "signal", required_argument, NULL, 's'},
908                 #ifdef LINUX
909                 { "monitor-forks", no_argument, NULL, 'm'},
910                 #endif
911                 { NULL, 0, NULL, 0 }
912         };
913         //argument variables
914         const char *exe=NULL;
915         const char *path=NULL;
916         int perclimit=0;
917         int pid_ok = FALSE;
918         int process_ok = FALSE;
919         int limit_ok = FALSE;
920         int last_known_argument = 0;
921         int kill_process = FALSE;   // kill process instead of stopping it
922         int restore_process = FALSE;  // restore killed process
923         int run_child_in_background = TRUE;  // run cpulimit in background when 
924                                              //  we launch new process
925         // struct rlimit maxlimit;
926
927         NCPU = get_ncpu();
928
929         opterr = 0;      // avoid unwanted error messages for unknown parameters
930         do {
931                 next_option = getopt_long (argc, argv, short_options,long_options, NULL);
932                 switch(next_option) {
933                         case 'b':
934                                 run_in_background = TRUE;
935                                 last_known_argument++;
936                                 break;
937                         case 'f':
938                                 run_child_in_background = FALSE;
939                                 run_in_background = FALSE;
940                                 last_known_argument++;
941                                 break;
942                         case 'p':
943                                 pid=atoi(optarg);
944                                 if (pid)   // valid PID
945                                 {
946                                   pid_ok = TRUE;
947                                   lazy = TRUE;
948                                 }
949                                 last_known_argument += 2;
950                                 break;
951                         case 'e':
952                                 exe=optarg;
953                                 process_ok = TRUE;
954                                 last_known_argument += 2;
955                                 break;
956                         case 'P':
957                                 path=optarg;
958                                 process_ok = TRUE;
959                                 last_known_argument += 2;
960                                 break;
961                         case 'l':
962                                 perclimit=atoi(optarg);
963                                 limit_ok = TRUE;
964                                 last_known_argument += 2;
965                                 break;
966                         case 'c':
967                                 NCPU = atoi(optarg);
968                                 last_known_argument += 2;
969                                 break;
970                         case 's':
971                                 send_signal = Translate_Signal(optarg);
972                                 if ( (send_signal < 1) || (send_signal > 35) )
973                                 {
974                                     fprintf(stderr, "Specified exit signal is not recognized or not within bounds (1-35). Using SIGCONT.\n");
975                                     send_signal = SIGCONT;
976                                 }
977                                 last_known_argument += 2;
978                         case 'k':
979                                 kill_process = TRUE;
980                                 last_known_argument++;
981                                 break;
982                         #ifdef LINUX
983                         case 'm':
984                                 monitor_children = TRUE;
985                                 last_known_argument++;
986                                 break;
987                         #endif
988                         case 'r':
989                                 restore_process = TRUE;
990                                 last_known_argument++;
991                                 break;
992
993                         case 'v':
994                                 verbose = TRUE;
995                                 last_known_argument++;
996                                 break;
997                         case 'q':
998                                 quiet = TRUE;
999                                 last_known_argument++;
1000                                 break;
1001                         case 'z':
1002                                 lazy = TRUE;
1003                                 last_known_argument++;
1004                                 break;
1005                         case 'h':
1006                                 print_usage (stdout, 1);
1007                                 last_known_argument++;
1008                                 break;
1009                         case 'o':
1010                                 last_known_argument++;
1011                                 next_option = -1;
1012                                 break;
1013                         case '?':
1014                                 print_usage (stderr, 1);
1015                                 last_known_argument++;
1016                                 break;
1017                         case -1:
1018                                 break;
1019                         // default:
1020                         //      abort();
1021                 }
1022         } while(next_option != -1);
1023
1024         signal(SIGCHLD, Child_Done);
1025
1026         // try to launch a program passed on the command line
1027         // But only if we do not already have a PID to watch
1028         if ( (last_known_argument + 1 < argc) && (pid_ok == FALSE) )
1029         {
1030            last_known_argument++;
1031            // if we stopped on "--" jump to the next parameter
1032            if ( (last_known_argument + 1 < argc) && (! strcmp(argv[last_known_argument], "--") ) )
1033                last_known_argument++;
1034            pid_t forked_pid;
1035            // try to launch remaining arguments
1036            if (verbose)
1037            {
1038                int index = last_known_argument;
1039                printf("Launching %s", argv[index]);
1040                for (index = last_known_argument + 1; index < argc; index++)
1041                     printf(" %s", argv[index]);
1042                printf(" with limit %d\n", perclimit);
1043            }
1044            forked_pid = fork();
1045            if (forked_pid == -1)  // error
1046            {
1047                printf("Failed to launch specified process.\n");
1048                exit(1);
1049            }
1050            else if (forked_pid == 0)   // target child
1051            {
1052               execvp(argv[last_known_argument],
1053                      &(argv[last_known_argument]) );
1054               exit(2);
1055            }
1056            else     // parent who will now fork the throttler
1057            {
1058               pid_t limit_pid;
1059               // if we are planning to kill a process, give it
1060               // a running head start to avoid death at start-up
1061               if (kill_process)
1062                  sleep(5);
1063     
1064               /* The following block assumes we want to run cpulimit in the
1065                  background. This is the default behaviour.
1066               */ 
1067               if (run_child_in_background)
1068               {
1069                  limit_pid = fork();
1070                  if (limit_pid == 0)   // child cpulimit process running in background
1071                  {
1072                     pid = forked_pid;    // the first child, target process
1073                     lazy = TRUE;
1074                     pid_ok = TRUE;
1075                     if (verbose)
1076                       printf("Throttling process %d\n", (int) pid);
1077                  }
1078                  else    // parent cpulimit process which can quit
1079                    exit(0);
1080               }  // end of running in background
1081               else
1082               {
1083                   pid = forked_pid;
1084                   lazy = TRUE;
1085                   pid_ok = TRUE;
1086                   run_in_background = FALSE;
1087               }  // end of running in foreground
1088
1089            }  // end of parent that launched target
1090
1091         }      // end of launching child process
1092
1093         if (!process_ok && !pid_ok) {
1094                 fprintf(stderr,"Error: You must specify a target process\n");
1095                 print_usage (stderr, 1);
1096                 exit(1);
1097         }
1098         if ((exe!=NULL && path!=NULL) || (pid_ok && (exe!=NULL || path!=NULL))) {
1099                 fprintf(stderr,"Error: You must specify exactly one target process\n");
1100                 print_usage (stderr, 1);
1101                 exit(1);
1102         }
1103         if (!limit_ok) {
1104                 fprintf(stderr,"Error: You must specify a cpu limit\n");
1105                 print_usage (stderr, 1);
1106                 exit(1);
1107         }
1108         float limit=perclimit/100.0;
1109         if ( (limit <= 0.00) || (limit > NCPU) )
1110         {
1111                 fprintf(stderr,"Error: limit must be in the range of 1 to %d00\n", NCPU);
1112                 print_usage (stderr, 1);
1113                 exit(1);
1114         }
1115
1116         // check to see if we should fork
1117         if (run_in_background)
1118         {
1119              pid_t process_id;
1120              process_id = fork();
1121              if (! process_id)
1122                 exit(0);
1123              else
1124              {
1125                 setsid();
1126                 process_id = fork();
1127                 if (process_id)
1128                   exit(0);
1129              }
1130         }
1131
1132         //parameters are all ok!
1133         signal(SIGINT,quit);
1134         signal(SIGTERM,quit);
1135
1136         my_pid = getpid();
1137         if (verbose)
1138            printf("%d CPUs detected.\n", NCPU);
1139
1140         increase_priority();
1141
1142         //time quantum in microseconds. it's splitted in a working period and a sleeping one
1143         int period=100000;
1144         struct timespec twork,tsleep;   //working and sleeping intervals
1145         memset(&twork,0,sizeof(struct timespec));
1146         memset(&tsleep,0,sizeof(struct timespec));
1147
1148 wait_for_process:
1149
1150         //look for the target process..or wait for it
1151         if (exe != NULL)
1152                 pid=getpidof(exe);
1153         else if (path != NULL)
1154                 pid=getpidof(path);
1155         else 
1156                 waitforpid(pid);
1157
1158         
1159         //process detected...let's play
1160
1161         //init compute_cpu_usage internal stuff
1162         compute_cpu_usage(0,0,NULL);
1163         //main loop counter
1164         int i=0;
1165
1166         struct timespec startwork,endwork;
1167         long workingtime=0;             //last working time in microseconds
1168
1169         if (verbose) print_caption();
1170
1171         float pcpu_avg=0;
1172
1173         // On Linux we can monitor child processes of the target
1174         #ifdef LINUX
1175         if (monitor_children)
1176         {
1177            pthread_t my_thread;
1178            int thread_status;
1179            if (verbose)
1180               printf("Starting fork monitoring thread...\n");
1181            program_data.this_program = argv[0];
1182            program_data.limit = perclimit;
1183            thread_status = pthread_create(&my_thread, NULL, 
1184                                           Monitor_Children, &program_data);
1185            if ( (thread_status) && (verbose) )
1186               printf("Creating fork monitoring thread failed.\n");
1187         }
1188         #endif
1189
1190         //here we should already have high priority, for time precision
1191         while(1) {
1192
1193                 //estimate how much the controlled process is using the cpu in its working interval
1194                 struct cpu_usage cu;
1195                 if (compute_cpu_usage(pid,workingtime,&cu)==-1) {
1196             if (!quiet)
1197                         fprintf(stderr,"Process %d dead!\n",pid);
1198                         if (lazy) exit(2);
1199                         //wait until our process appears
1200                         goto wait_for_process;          
1201                 }
1202
1203                 //cpu actual usage of process (range 0-1)
1204                 float pcpu=cu.pcpu;
1205                 //rate at which we are keeping active the process (range 0-1)
1206                 float workingrate=cu.workingrate;
1207
1208                 //adjust work and sleep time slices
1209                 if (pcpu>0) {
1210                         twork.tv_nsec=min(period*limit*1000/pcpu*workingrate,period*1000);
1211                 }
1212                 else if (pcpu==0) {
1213                         twork.tv_nsec=period*1000;
1214                 }
1215                 else if (pcpu==-1) {
1216                         //not yet a valid idea of cpu usage
1217                         pcpu=limit;
1218                         workingrate=limit;
1219                         twork.tv_nsec=min(period*limit*1000,period*1000);
1220                 }
1221                 tsleep.tv_nsec=period*1000-twork.tv_nsec;
1222
1223                 //update average usage
1224                 pcpu_avg=(pcpu_avg*i+pcpu)/(i+1);
1225
1226                 if (verbose && i%10==0 && i>0) {
1227                         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);
1228                         if (i%200 == 0)
1229                            print_caption();
1230                 }
1231
1232                 // if (limit<1 && limit>0) {
1233                 // printf("Comparing %f to %f\n", pcpu, limit);
1234                 if (pcpu < limit)
1235                 {
1236                         // printf("Continue\n");
1237                         //resume process
1238                         if (kill(pid,SIGCONT)!=0) {
1239                              if (!quiet)
1240                                 fprintf(stderr,"Process %d dead!\n",pid);
1241                              if (lazy) exit(2);
1242                                 //wait until our process appears
1243                                 goto wait_for_process;
1244                         }
1245                 }
1246
1247                 #ifdef __APPLE_
1248                 // OS X does not have clock_gettime, use clock_get_time
1249                 clock_serv_t cclock;
1250                 mach_timespec_t mts;
1251                 host_get_clock_service(mach_host_self(), CALENDAR_CLOCK, &cclock);
1252                 clock_get_time(cclock, &mts);
1253                 mach_port_deallocate(mach_task_self(), cclock);
1254                 startwork.tv_sec = mts.tv_sec;
1255                 startwork.tv_nsec = mts.tv_nsec;
1256
1257                 #else
1258                 clock_gettime(CLOCK_REALTIME,&startwork);
1259                 #endif
1260
1261                 nanosleep(&twork,NULL);         //now process is working
1262                 #ifdef __APPLE__
1263                 // OS X does not have clock_gettime, use clock_get_time
1264                 // clock_serv_t cclock;
1265                 // mach_timespec_t mts;
1266                 host_get_clock_service(mach_host_self(), CALENDAR_CLOCK, &cclock);
1267                 clock_get_time(cclock, &mts);
1268                 mach_port_deallocate(mach_task_self(), cclock);
1269                 endwork.tv_sec = mts.tv_sec;
1270                 endwork.tv_nsec = mts.tv_nsec;
1271
1272                 #else
1273                 clock_gettime(CLOCK_REALTIME,&endwork);
1274                 #endif
1275                 workingtime=timediff(&endwork,&startwork);
1276
1277                 // if (limit<1) {
1278                 // printf("Checking %f vs %f\n", pcpu, limit);
1279                 if (pcpu > limit)
1280                 {
1281                      // When over our limit we may run into
1282                      // situations where we want to kill
1283                      // the offending process, then restart it
1284                      if (kill_process)
1285                      {
1286                          kill(pid, SIGKILL);
1287                          if (!quiet)
1288                              fprintf(stderr, "Process %d killed.\n", pid);
1289                          if ( (lazy) && (! restore_process) ) 
1290                               exit(2);
1291                          // restart killed process
1292                          if (restore_process)
1293                          {
1294                              pid_t new_process;
1295                              new_process = fork();
1296                              if (new_process == -1)
1297                              {
1298                               fprintf(stderr, "Failed to restore killed process.\n");
1299                              }
1300                              else if (new_process == 0)
1301                              {
1302                                 // child which becomes new process
1303                                 if (verbose)
1304                                    printf("Relaunching %s\n",
1305                                           argv[last_known_argument]);
1306                                 execvp(argv[last_known_argument],
1307                                        &(argv[last_known_argument]) ); 
1308                              }
1309                              else // parent
1310                              {
1311                                 // we need to track new process
1312                                 pid = new_process;
1313                                 // avoid killing child process
1314                                 sleep(5);
1315                              }
1316                          }
1317                      }
1318                      // do not kill process, just throttle it
1319                      else
1320                      {
1321
1322                         // printf("Stop\n");
1323                         //stop process, it has worked enough
1324                         if (kill(pid,SIGSTOP)!=0) {
1325                             if (!quiet)
1326                                 fprintf(stderr,"Process %d dead!\n", pid);
1327                             if (lazy) exit(2);
1328                                 //wait until our process appears
1329                                 goto wait_for_process;
1330                         }
1331                         nanosleep(&tsleep,NULL);        //now process is sleeping
1332                       }   // end of throttle process
1333                 }         // end of process using too much CPU
1334                 i++;
1335         }
1336
1337    return 0;
1338 }