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