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  *
  19  */
  20
  21
  22 #include <getopt.h>
  23 #include <stdio.h>
  24 #include <stdlib.h>
  25 #include <time.h>
  26 #include <sys/time.h>
  27 #include <unistd.h>
  28 #include <sys/types.h>
  29 #include <signal.h>
  30 #include <sys/resource.h>
  31 #include <string.h>
  32 #include <dirent.h>
  33 #include <errno.h>
  34 #include <string.h>
  35 #include <limits.h>    // for compatibility
  36
  37
  38 //kernel time resolution (inverse of one jiffy interval) in Hertz
  39 //i don't know how to detect it, then define to the default (not very clean!)
  40 #define HZ 100
  41
  42 //some useful macro
  43 #define min(a,b) (a<b?a:b)
  44 #define max(a,b) (a>b?a:b)
  45
  46 //pid of the controlled process
  47 int pid=0;
  48 //executable file name
  49 char *program_name;
  50 //verbose mode
  51 int verbose=0;
  52 //lazy mode
  53 int lazy=0;
  54 // is higher priority nice possible?
  55 int nice_lim;
  56
  57 //reverse byte search
  58 void *memrchr(const void *s, int c, size_t n);
  59
  60 //return ta-tb in microseconds (no overflow checks!)
  61 inline long timediff(const struct timespec *ta,const struct timespec *tb) {
  62     unsigned long us = (ta->tv_sec-tb->tv_sec)*1000000 + (ta->tv_nsec/1000 - tb->tv_nsec/1000);
  63     return us;
  64 }
  65
  66 int waitforpid(int pid) {
  67         //switch to low priority
  68         // if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
  69         if ( (nice_lim < INT_MAX) &&
  70              (setpriority(PRIO_PROCESS, getpid(), 19) != 0) ) {
  71                 printf("Warning: cannot renice\n");
  72         }
  73
  74         int i=0;
  75
  76         while(1) {
  77
  78                 DIR *dip;
  79                 struct dirent *dit;
  80
  81                 //open a directory stream to /proc directory
  82                 if ((dip = opendir("/proc")) == NULL) {
  83                         perror("opendir");
  84                         return -1;
  85                 }
  86
  87                 //read in from /proc and seek for process dirs
  88                 while ((dit = readdir(dip)) != NULL) {
  89                         //get pid
  90                         if (pid==atoi(dit->d_name)) {
  91                                 //pid detected
  92                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
  93                                         //process is ok!
  94                                         if (closedir(dip) == -1) {
  95                                            perror("closedir");
  96                                            return -1;
  97                                         }
  98                                         goto done;
  99                                 }
 100                                 else {
 101                                         fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
 102                                 }
 103                         }
 104                 }
 105
 106                 //close the dir stream and check for errors
 107                 if (closedir(dip) == -1) {
 108                         perror("closedir");
 109                         return -1;
 110                 }
 111
 112                 //no suitable target found
 113                 if (i++==0) {
 114                         if (lazy) {
 115                                 fprintf(stderr,"No process found\n");
 116                                 exit(2);
 117                         }
 118                         else {
 119                                 printf("Warning: no target process found. Waiting for it...\n");
 120                         }
 121                 }
 122
 123                 //sleep for a while
 124                 sleep(2);
 125         }
 126
 127 done:
 128         printf("Process %d detected\n",pid);
 129         //now set high priority, if possible
 130         // if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
 131         if ( (nice_lim < INT_MAX) &&
 132              (setpriority(PRIO_PROCESS, getpid(), nice_lim) != 0) ) {
 133                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
 134         }
 135         return 0;
 136
 137 }
 138
 139 //this function periodically scans process list and looks for executable path names
 140 //it should be executed in a low priority context, since precise timing does not matter
 141 //if a process is found then its pid is returned
 142 //process: the name of the wanted process, can be an absolute path name to the executable file
 143 //         or simply its name
 144 //return: pid of the found process
 145 int getpidof(const char *process) {
 146
 147         //set low priority
 148         // if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
 149         if ( (nice_lim < INT_MAX) &&
 150              (setpriority(PRIO_PROCESS, getpid(), 19) != 0) ) {
 151                 printf("Warning: cannot renice\n");
 152         }
 153
 154         char exelink[20];
 155         char exepath[PATH_MAX+1];
 156         int pid=0;
 157         int i=0;
 158
 159         while(1) {
 160
 161                 DIR *dip;
 162                 struct dirent *dit;
 163
 164                 //open a directory stream to /proc directory
 165                 if ((dip = opendir("/proc")) == NULL) {
 166                         perror("opendir");
 167                         return -1;
 168                 }
 169
 170                 //read in from /proc and seek for process dirs
 171                 while ((dit = readdir(dip)) != NULL) {
 172                         //get pid
 173                         pid=atoi(dit->d_name);
 174                         if (pid>0) {
 175                                 sprintf(exelink,"/proc/%d/exe",pid);
 176                                 int size=readlink(exelink,exepath,sizeof(exepath));
 177                                 if (size>0) {
 178                                         int found=0;
 179                                         if (process[0]=='/' && strncmp(exepath,process,size)==0 && size==strlen(process)) {
 180                                                 //process starts with / then it's an absolute path
 181                                                 found=1;
 182                                         }
 183                                         else {
 184                                                 //process is the name of the executable file
 185                                                 if (strncmp(exepath+size-strlen(process),process,strlen(process))==0) {
 186                                                         found=1;
 187                                                 }
 188                                         }
 189                                         if (found==1) {
 190                                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
 191                                                         //process is ok!
 192                                                         if (closedir(dip) == -1) {
 193                                                           perror("closedir");
 194                                                           return -1;
 195                                                         }
 196                                                         goto done;
 197                                                 }
 198                                                 else {
 199                                                         fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
 200                                                 }
 201                                         }
 202                                 }
 203                         }
 204                 }
 205
 206                 //close the dir stream and check for errors
 207                 if (closedir(dip) == -1) {
 208                         perror("closedir");
 209                         return -1;
 210                 }
 211
 212                 //no suitable target found
 213                 if (i++==0) {
 214                         if (lazy) {
 215                                 fprintf(stderr,"No process found\n");
 216                                 exit(2);
 217                         }
 218                         else {
 219                                 printf("Warning: no target process found. Waiting for it...\n");
 220                         }
 221                 }
 222
 223                 //sleep for a while
 224                 sleep(2);
 225         }
 226
 227 done:
 228         printf("Process %d detected\n",pid);
 229         //now set high priority, if possible
 230         // if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
 231         if ( (nice_lim < INT_MAX) &&
 232              (setpriority(PRIO_PROCESS, getpid(), nice_lim) != 0) ) {
 233                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
 234         }
 235         return pid;
 236
 237 }
 238
 239 //SIGINT and SIGTERM signal handler
 240 void quit(int sig) {
 241         //let the process continue if it's stopped
 242         kill(pid,SIGCONT);
 243         printf("Exiting...\n");
 244         exit(0);
 245 }
 246
 247 //get jiffies count from /proc filesystem
 248 int getjiffies(int pid) {
 249         static char stat[20];
 250         static char buffer[1024];
 251         char *p;
 252         sprintf(stat,"/proc/%d/stat",pid);
 253         FILE *f=fopen(stat,"r");
 254         if (f==NULL) return -1;
 255         p = fgets(buffer,sizeof(buffer),f);
 256         fclose(f);
 257         // char *p=buffer;
 258         if (p)
 259         {
 260           p=memchr(p+1,')',sizeof(buffer)-(p-buffer));
 261           int sp=12;
 262           while (sp--)
 263                 p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
 264           //user mode jiffies
 265           int utime=atoi(p+1);
 266           p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
 267           //kernel mode jiffies
 268           int ktime=atoi(p+1);
 269           return utime+ktime;
 270         }
 271         // could not read info
 272         return -1;
 273 }
 274
 275 //process instant photo
 276 struct process_screenshot {
 277         struct timespec when;   //timestamp
 278         int jiffies;    //jiffies count of the process
 279         int cputime;    //microseconds of work from previous screenshot to current
 280 };
 281
 282 //extracted process statistics
 283 struct cpu_usage {
 284         float pcpu;
 285         float workingrate;
 286 };
 287
 288 //this function is an autonomous dynamic system
 289 //it works with static variables (state variables of the system), that keep memory of recent past
 290 //its aim is to estimate the cpu usage of the process
 291 //to work properly it should be called in a fixed periodic way
 292 //perhaps i will put it in a separate thread...
 293 int compute_cpu_usage(int pid,int last_working_quantum,struct cpu_usage *pusage) {
 294         #define MEM_ORDER 10
 295         //circular buffer containing last MEM_ORDER process screenshots
 296         static struct process_screenshot ps[MEM_ORDER];
 297         //the last screenshot recorded in the buffer
 298         static int front=-1;
 299         //the oldest screenshot recorded in the buffer
 300         static int tail=0;
 301
 302         if (pusage==NULL) {
 303                 //reinit static variables
 304                 front=-1;
 305                 tail=0;
 306                 return 0;
 307         }
 308
 309         //let's advance front index and save the screenshot
 310         front=(front+1)%MEM_ORDER;
 311         int j=getjiffies(pid);
 312         if (j>=0) ps[front].jiffies=j;
 313         else return -1; //error: pid does not exist
 314         clock_gettime(CLOCK_REALTIME,&(ps[front].when));
 315         ps[front].cputime=last_working_quantum;
 316
 317         //buffer actual size is: (front-tail+MEM_ORDER)%MEM_ORDER+1
 318         int size=(front-tail+MEM_ORDER)%MEM_ORDER+1;
 319
 320         if (size==1) {
 321                 //not enough samples taken (it's the first one!), return -1
 322                 pusage->pcpu=-1;
 323                 pusage->workingrate=1;
 324                 return 0;
 325         }
 326         else {
 327                 //now we can calculate cpu usage, interval dt and dtwork are expressed in microseconds
 328                 long dt=timediff(&(ps[front].when),&(ps[tail].when));
 329                 long dtwork=0;
 330                 int i=(tail+1)%MEM_ORDER;
 331                 int max=(front+1)%MEM_ORDER;
 332                 do {
 333                         dtwork+=ps[i].cputime;
 334                         i=(i+1)%MEM_ORDER;
 335                 } while (i!=max);
 336                 int used=ps[front].jiffies-ps[tail].jiffies;
 337                 float usage=(used*1000000.0/HZ)/dtwork;
 338                 pusage->workingrate=1.0*dtwork/dt;
 339                 pusage->pcpu=usage*pusage->workingrate;
 340                 if (size==MEM_ORDER)
 341                         tail=(tail+1)%MEM_ORDER;
 342                 return 0;
 343         }
 344         #undef MEM_ORDER
 345 }
 346
 347 void print_caption() {
 348         printf("\n%%CPU\twork quantum\tsleep quantum\tactive rate\n");
 349 }
 350
 351 void print_usage(FILE *stream,int exit_code) {
 352         fprintf(stream, "Usage: %s TARGET [OPTIONS...]\n",program_name);
 353         fprintf(stream, "   TARGET must be exactly one of these:\n");
 354         fprintf(stream, "      -p, --pid=N        pid of the process\n");
 355         fprintf(stream, "      -e, --exe=FILE     name of the executable program file\n");
 356         fprintf(stream, "      -P, --path=PATH    absolute path name of the executable program file\n");
 357         fprintf(stream, "   OPTIONS\n");
 358         fprintf(stream, "      -b  --background   run in background\n");
 359         fprintf(stream, "      -l, --limit=N      percentage of cpu allowed from 0 to 100 (mandatory)\n");
 360         fprintf(stream, "      -v, --verbose      show control statistics\n");
 361         fprintf(stream, "      -z, --lazy         exit if there is no suitable target process, or if it dies\n");
 362         fprintf(stream, "      -h, --help         display this help and exit\n");
 363         exit(exit_code);
 364 }
 365
 366 int main(int argc, char **argv) {
 367
 368         //get program name
 369         char *p=(char*)memrchr(argv[0],(unsigned int)'/',strlen(argv[0]));
 370         program_name = p==NULL?argv[0]:(p+1);
 371         int run_in_background = 0;
 372         //parse arguments
 373         int next_option;
 374         /* A string listing valid short options letters. */
 375         const char* short_options="p:e:P:l:bvzh";
 376         /* An array describing valid long options. */
 377         const struct option long_options[] = {
 378                 { "pid", required_argument, NULL, 'p' },
 379                 { "exe", required_argument, NULL, 'e' },
 380                 { "path", required_argument, NULL, 'P' },
 381                 { "limit", required_argument, NULL, 'l' },
 382                 { "background", no_argument, NULL, 'b' },
 383                 { "verbose", no_argument, NULL, 'v' },
 384                 { "lazy", no_argument, NULL, 'z' },
 385                 { "help", no_argument, NULL, 'h' },
 386                 { NULL, 0, NULL, 0 }
 387         };
 388         //argument variables
 389         const char *exe=NULL;
 390         const char *path=NULL;
 391         int perclimit=0;
 392         int pid_ok=0;
 393         int process_ok=0;
 394         int limit_ok=0;
 395         struct rlimit maxlimit;
 396
 397         do {
 398                 next_option = getopt_long (argc, argv, short_options,long_options, NULL);
 399                 switch(next_option) {
 400                         case 'b':
 401                                 run_in_background = 1;
 402                                 break;
 403                         case 'p':
 404                                 pid=atoi(optarg);
 405                                 pid_ok=1;
 406                                 lazy = 1;
 407                                 break;
 408                         case 'e':
 409                                 exe=optarg;
 410                                 process_ok=1;
 411                                 break;
 412                         case 'P':
 413                                 path=optarg;
 414                                 process_ok=1;
 415                                 break;
 416                         case 'l':
 417                                 perclimit=atoi(optarg);
 418                                 limit_ok=1;
 419                                 break;
 420                         case 'v':
 421                                 verbose=1;
 422                                 break;
 423                         case 'z':
 424                                 lazy=1;
 425                                 break;
 426                         case 'h':
 427                                 print_usage (stdout, 1);
 428                                 break;
 429                         case '?':
 430                                 print_usage (stderr, 1);
 431                                 break;
 432                         case -1:
 433                                 break;
 434                         default:
 435                                 abort();
 436                 }
 437         } while(next_option != -1);
 438
 439         if (!process_ok && !pid_ok) {
 440                 fprintf(stderr,"Error: You must specify a target process\n");
 441                 print_usage (stderr, 1);
 442                 exit(1);
 443         }
 444         if ((exe!=NULL && path!=NULL) || (pid_ok && (exe!=NULL || path!=NULL))) {
 445                 fprintf(stderr,"Error: You must specify exactly one target process\n");
 446                 print_usage (stderr, 1);
 447                 exit(1);
 448         }
 449         if (!limit_ok) {
 450                 fprintf(stderr,"Error: You must specify a cpu limit\n");
 451                 print_usage (stderr, 1);
 452                 exit(1);
 453         }
 454         float limit=perclimit/100.0;
 455         if (limit <= 0.00) // || limit >1) {
 456         {
 457                 fprintf(stderr,"Error: limit must be greater than 0\n");
 458                 print_usage (stderr, 1);
 459                 exit(1);
 460         }
 461
 462         // check to see if we should fork
 463         if (run_in_background)
 464         {
 465              pid_t process_id;
 466              process_id = fork();
 467              if (! process_id)
 468                 exit(0);
 469              else
 470              {
 471                 setsid();
 472                 process_id = fork();
 473                 if (process_id)
 474                   exit(0);
 475              }
 476         }
 477
 478         //parameters are all ok!
 479         signal(SIGINT,quit);
 480         signal(SIGTERM,quit);
 481
 482         if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
 483         //if that failed, check if we have a limit
 484         // by how much we can raise the priority
 485 #ifdef RLIMIT_NICE
 486 //check if non-root can even make changes
 487 // (ifdef because it's only available in linux >= 2.6.13)
 488                 nice_lim=getpriority(PRIO_PROCESS,getpid());
 489                 getrlimit(RLIMIT_NICE, &maxlimit);
 490
 491 //if we can do better then current
 492                 if( (20 - (signed)maxlimit.rlim_cur) < nice_lim &&
 493                     setpriority(PRIO_PROCESS,getpid(),
 494                     20 - (signed)maxlimit.rlim_cur)==0 //and it actually works
 495                   ) {
 496
 497                         //if we can do better, but not by much, warn about it
 498                         if( (nice_lim - (20 - (signed)maxlimit.rlim_cur)) < 9)
 499                         {
 500                         printf("Warning, can only increase priority by %d.\n",                                nice_lim - (20 - (signed)maxlimit.rlim_cur));
 501                         }
 502                         //our new limit
 503                         nice_lim = 20 - (signed)maxlimit.rlim_cur;
 504
 505                 } else
 506 // otherwise don't try to change priority.
 507 // The below will also run if it's not possible
 508 // for non-root to change priority
 509 #endif
 510                 {
 511                         printf("Warning: cannot renice.\nTo work better you should run this program as root, or adjust RLIMIT_NICE.\nFor example in /etc/security/limits.conf add a line with: * - nice -10\n\n");
 512                         nice_lim=INT_MAX;
 513                 }
 514         } else {
 515                 nice_lim=-20;
 516         }
 517         //don't bother putting setpriority back down,
 518         // since getpidof and waitforpid twiddle it anyway
 519
 520
 521
 522         //time quantum in microseconds. it's splitted in a working period and a sleeping one
 523         int period=100000;
 524         struct timespec twork,tsleep;   //working and sleeping intervals
 525         memset(&twork,0,sizeof(struct timespec));
 526         memset(&tsleep,0,sizeof(struct timespec));
 527
 528 wait_for_process:
 529
 530         //look for the target process..or wait for it
 531         if (exe!=NULL)
 532                 pid=getpidof(exe);
 533         else if (path!=NULL)
 534                 pid=getpidof(path);
 535         else {
 536                 waitforpid(pid);
 537         }
 538         //process detected...let's play
 539
 540         //init compute_cpu_usage internal stuff
 541         compute_cpu_usage(0,0,NULL);
 542         //main loop counter
 543         int i=0;
 544
 545         struct timespec startwork,endwork;
 546         long workingtime=0;             //last working time in microseconds
 547
 548         if (verbose) print_caption();
 549
 550         float pcpu_avg=0;
 551
 552         //here we should already have high priority, for time precision
 553         while(1) {
 554
 555                 //estimate how much the controlled process is using the cpu in its working interval
 556                 struct cpu_usage cu;
 557                 if (compute_cpu_usage(pid,workingtime,&cu)==-1) {
 558                         fprintf(stderr,"Process %d dead!\n",pid);
 559                         if (lazy) exit(2);
 560                         //wait until our process appears
 561                         goto wait_for_process;
 562                 }
 563
 564                 //cpu actual usage of process (range 0-1)
 565                 float pcpu=cu.pcpu;
 566                 //rate at which we are keeping active the process (range 0-1)
 567                 float workingrate=cu.workingrate;
 568
 569                 //adjust work and sleep time slices
 570                 if (pcpu>0) {
 571                         twork.tv_nsec=min(period*limit*1000/pcpu*workingrate,period*1000);
 572                 }
 573                 else if (pcpu==0) {
 574                         twork.tv_nsec=period*1000;
 575                 }
 576                 else if (pcpu==-1) {
 577                         //not yet a valid idea of cpu usage
 578                         pcpu=limit;
 579                         workingrate=limit;
 580                         twork.tv_nsec=min(period*limit*1000,period*1000);
 581                 }
 582                 tsleep.tv_nsec=period*1000-twork.tv_nsec;
 583
 584                 //update average usage
 585                 pcpu_avg=(pcpu_avg*i+pcpu)/(i+1);
 586
 587                 if (verbose && i%10==0 && i>0) {
 588                         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);
 589                 }
 590
 591                 if (limit<1 && limit>0) {
 592                         //resume process
 593                         if (kill(pid,SIGCONT)!=0) {
 594                                 fprintf(stderr,"Process %d dead!\n",pid);
 595                                 if (lazy) exit(2);
 596                                 //wait until our process appears
 597                                 goto wait_for_process;
 598                         }
 599                 }
 600
 601                 clock_gettime(CLOCK_REALTIME,&startwork);
 602                 nanosleep(&twork,NULL);         //now process is working
 603                 clock_gettime(CLOCK_REALTIME,&endwork);
 604                 workingtime=timediff(&endwork,&startwork);
 605
 606                 if (limit<1) {
 607                         //stop process, it has worked enough
 608                         if (kill(pid,SIGSTOP)!=0) {
 609                                 fprintf(stderr,"Process %d dead!\n",pid);
 610                                 if (lazy) exit(2);
 611                                 //wait until our process appears
 612                                 goto wait_for_process;
 613                         }
 614                         nanosleep(&tsleep,NULL);        //now process is sleeping
 615                 }
 616                 i++;
 617         }
 618
 619 }