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