cpulimit.c

   1 /**
   2  * Copyright (c) 2005, by:      Angelo Marletta <marlonx80@hotmail.com>
   3  *
   4  * This file may be used subject to the terms and conditions of the
   5  * GNU Library General Public License Version 2, or any later version
   6  * at your option, as published by the Free Software Foundation.
   7  * This program is distributed in the hope that it will be useful,
   8  * but WITHOUT ANY WARRANTY; without even the implied warranty of
   9  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  10  * GNU Library General Public License for more details.
  11  *
  12  **********************************************************************
  13  *
  14  * Simple program to limit the cpu usage of a process
  15  * If you modify this code, send me a copy please
  16  *
  17  * Author:  Angelo Marletta
  18  * Date:    26/06/2005
  19  * Version: 1.1
  20  * Last version at: http://marlon80.interfree.it/cpulimit/index.html
  21  *
  22  * Changelog:
  23  *  - Fixed a segmentation fault if controlled process exited in particular circumstances
  24  *  - Better CPU usage estimate
  25  *  - Fixed a <0 %CPU usage reporting in rare cases
  26  *  - Replaced MAX_PATH_SIZE with PATH_MAX already defined in <limits.h>
  27  *  - Command line arguments now available
  28  *  - Now is possible to specify target process by pid
  29  */
  30
  31
  32 #include <getopt.h>
  33 #include <stdio.h>
  34 #include <stdlib.h>
  35 #include <time.h>
  36 #include <sys/time.h>
  37 #include <unistd.h>
  38 #include <sys/types.h>
  39 #include <signal.h>
  40 #include <sys/resource.h>
  41 #include <string.h>
  42 #include <dirent.h>
  43 #include <errno.h>
  44 #include <string.h>
  45
  46 //kernel time resolution (inverse of one jiffy interval) in Hertz
  47 //i don't know how to detect it, then define to the default (not very clean!)
  48 #define HZ 100
  49
  50 //some useful macro
  51 #define min(a,b) (a<b?a:b)
  52 #define max(a,b) (a>b?a:b)
  53
  54 //pid of the controlled process
  55 int pid=0;
  56 //executable file name
  57 char *program_name;
  58 //verbose mode
  59 int verbose=0;
  60 //lazy mode
  61 int lazy=0;
  62
  63 //reverse byte search
  64 void *memrchr(const void *s, int c, size_t n);
  65
  66 //return ta-tb in microseconds (no overflow checks!)
  67 inline long timediff(const struct timespec *ta,const struct timespec *tb) {
  68     unsigned long us = (ta->tv_sec-tb->tv_sec)*1000000 + (ta->tv_nsec/1000 - tb->tv_nsec/1000);
  69     return us;
  70 }
  71
  72 int waitforpid(int pid) {
  73         //switch to low priority
  74         if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
  75                 printf("Warning: cannot renice\n");
  76         }
  77
  78         int i=0;
  79
  80         while(1) {
  81
  82                 DIR *dip;
  83                 struct dirent *dit;
  84
  85                 //open a directory stream to /proc directory
  86                 if ((dip = opendir("/proc")) == NULL) {
  87                         perror("opendir");
  88                         return -1;
  89                 }
  90
  91                 //read in from /proc and seek for process dirs
  92                 while ((dit = readdir(dip)) != NULL) {
  93                         //get pid
  94                         if (pid==atoi(dit->d_name)) {
  95                                 //pid detected
  96                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
  97                                         //process is ok!
  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                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
 132         }
 133         return 0;
 134
 135 }
 136
 137 //this function periodically scans process list and looks for executable path names
 138 //it should be executed in a low priority context, since precise timing does not matter
 139 //if a process is found then its pid is returned
 140 //process: the name of the wanted process, can be an absolute path name to the executable file
 141 //         or simply its name
 142 //return: pid of the found process
 143 int getpidof(const char *process) {
 144
 145         //set low priority
 146         if (setpriority(PRIO_PROCESS,getpid(),19)!=0) {
 147                 printf("Warning: cannot renice\n");
 148         }
 149
 150         char exelink[20];
 151         char exepath[PATH_MAX+1];
 152         int pid=0;
 153         int i=0;
 154
 155         while(1) {
 156
 157                 DIR *dip;
 158                 struct dirent *dit;
 159
 160                 //open a directory stream to /proc directory
 161                 if ((dip = opendir("/proc")) == NULL) {
 162                         perror("opendir");
 163                         return -1;
 164                 }
 165
 166                 //read in from /proc and seek for process dirs
 167                 while ((dit = readdir(dip)) != NULL) {
 168                         //get pid
 169                         pid=atoi(dit->d_name);
 170                         if (pid>0) {
 171                                 sprintf(exelink,"/proc/%d/exe",pid);
 172                                 int size=readlink(exelink,exepath,sizeof(exepath));
 173                                 if (size>0) {
 174                                         int found=0;
 175                                         if (process[0]=='/' && strncmp(exepath,process,size)==0 && size==strlen(process)) {
 176                                                 //process starts with / then it's an absolute path
 177                                                 found=1;
 178                                         }
 179                                         else {
 180                                                 //process is the name of the executable file
 181                                                 if (strncmp(exepath+size-strlen(process),process,strlen(process))==0) {
 182                                                         found=1;
 183                                                 }
 184                                         }
 185                                         if (found==1) {
 186                                                 if (kill(pid,SIGSTOP)==0 &&  kill(pid,SIGCONT)==0) {
 187                                                         //process is ok!
 188                                                         goto done;
 189                                                 }
 190                                                 else {
 191                                                         fprintf(stderr,"Error: Process %d detected, but you don't have permission to control it\n",pid);
 192                                                 }
 193                                         }
 194                                 }
 195                         }
 196                 }
 197
 198                 //close the dir stream and check for errors
 199                 if (closedir(dip) == -1) {
 200                         perror("closedir");
 201                         return -1;
 202                 }
 203
 204                 //no suitable target found
 205                 if (i++==0) {
 206                         if (lazy) {
 207                                 fprintf(stderr,"No process found\n");
 208                                 exit(2);
 209                         }
 210                         else {
 211                                 printf("Warning: no target process found. Waiting for it...\n");
 212                         }
 213                 }
 214
 215                 //sleep for a while
 216                 sleep(2);
 217         }
 218
 219 done:
 220         printf("Process %d detected\n",pid);
 221         //now set high priority, if possible
 222         if (setpriority(PRIO_PROCESS,getpid(),-20)!=0) {
 223                 printf("Warning: cannot renice.\nTo work better you should run this program as root.\n");
 224         }
 225         return pid;
 226
 227 }
 228
 229 //SIGINT and SIGTERM signal handler
 230 void quit(int sig) {
 231         //let the process continue if it's stopped
 232         kill(pid,SIGCONT);
 233         printf("Exiting...\n");
 234         exit(0);
 235 }
 236
 237 //get jiffies count from /proc filesystem
 238 int getjiffies(int pid) {
 239         static char stat[20];
 240         static char buffer[1024];
 241         sprintf(stat,"/proc/%d/stat",pid);
 242         FILE *f=fopen(stat,"r");
 243         if (f==NULL) return -1;
 244         fgets(buffer,sizeof(buffer),f);
 245         fclose(f);
 246         char *p=buffer;
 247         p=memchr(p+1,')',sizeof(buffer)-(p-buffer));
 248         int sp=12;
 249         while (sp--)
 250                 p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
 251         //user mode jiffies
 252         int utime=atoi(p+1);
 253         p=memchr(p+1,' ',sizeof(buffer)-(p-buffer));
 254         //kernel mode jiffies
 255         int ktime=atoi(p+1);
 256         return utime+ktime;
 257 }
 258
 259 //process instant photo
 260 struct process_screenshot {
 261         struct timespec when;   //timestamp
 262         int jiffies;    //jiffies count of the process
 263         int cputime;    //microseconds of work from previous screenshot to current
 264 };
 265
 266 //extracted process statistics
 267 struct cpu_usage {
 268         float pcpu;
 269         float workingrate;
 270 };
 271
 272 //this function is an autonomous dynamic system
 273 //it works with static variables (state variables of the system), that keep memory of recent past
 274 //its aim is to estimate the cpu usage of the process
 275 //to work properly it should be called in a fixed periodic way
 276 //perhaps i will put it in a separate thread...
 277 int compute_cpu_usage(int pid,int last_working_quantum,struct cpu_usage *pusage) {
 278         #define MEM_ORDER 10
 279         //circular buffer containing last MEM_ORDER process screenshots
 280         static struct process_screenshot ps[MEM_ORDER];
 281         //the last screenshot recorded in the buffer
 282         static int front=-1;
 283         //the oldest screenshot recorded in the buffer
 284         static int tail=0;
 285
 286         if (pusage==NULL) {
 287                 //reinit static variables
 288                 front=-1;
 289                 tail=0;
 290                 return 0;
 291         }
 292
 293         //let's advance front index and save the screenshot
 294         front=(front+1)%MEM_ORDER;
 295         int j=getjiffies(pid);
 296         if (j>=0) ps[front].jiffies=j;
 297         else return -1; //error: pid does not exist
 298         clock_gettime(CLOCK_REALTIME,&(ps[front].when));
 299         ps[front].cputime=last_working_quantum;
 300
 301         //buffer actual size is: (front-tail+MEM_ORDER)%MEM_ORDER+1
 302         int size=(front-tail+MEM_ORDER)%MEM_ORDER+1;
 303
 304         if (size==1) {
 305                 //not enough samples taken (it's the first one!), return -1
 306                 pusage->pcpu=-1;
 307                 pusage->workingrate=1;
 308                 return 0;
 309         }
 310         else {
 311                 //now we can calculate cpu usage, interval dt and dtwork are expressed in microseconds
 312                 long dt=timediff(&(ps[front].when),&(ps[tail].when));
 313                 long dtwork=0;
 314                 int i=(tail+1)%MEM_ORDER;
 315                 int max=(front+1)%MEM_ORDER;
 316                 do {
 317                         dtwork+=ps[i].cputime;
 318                         i=(i+1)%MEM_ORDER;
 319                 } while (i!=max);
 320                 int used=ps[front].jiffies-ps[tail].jiffies;
 321                 float usage=(used*1000000.0/HZ)/dtwork;
 322                 pusage->workingrate=1.0*dtwork/dt;
 323                 pusage->pcpu=usage*pusage->workingrate;
 324                 if (size==MEM_ORDER)
 325                         tail=(tail+1)%MEM_ORDER;
 326                 return 0;
 327         }
 328         #undef MEM_ORDER
 329 }
 330
 331 void print_caption() {
 332         printf("\n%%CPU\twork quantum\tsleep quantum\tactive rate\n");
 333 }
 334
 335 void print_usage(FILE *stream,int exit_code) {
 336         fprintf(stream, "Usage: %s TARGET [OPTIONS...]\n",program_name);
 337         fprintf(stream, "   TARGET must be exactly one of these:\n");
 338         fprintf(stream, "      -p, --pid=N        pid of the process\n");
 339         fprintf(stream, "      -e, --exe=FILE     name of the executable program file\n");
 340         fprintf(stream, "      -P, --path=PATH    absolute path name of the executable program file\n");
 341         fprintf(stream, "   OPTIONS\n");
 342         fprintf(stream, "      -l, --limit=N      percentage of cpu allowed from 0 to 100 (mandatory)\n");
 343         fprintf(stream, "      -v, --verbose      show control statistics\n");
 344         fprintf(stream, "      -z, --lazy         exit if there is no suitable target process, or if it dies\n");
 345         fprintf(stream, "      -h, --help         display this help and exit\n");
 346         exit(exit_code);
 347 }
 348
 349 int main(int argc, char **argv) {
 350
 351         //get program name
 352         char *p=(char*)memrchr(argv[0],(unsigned int)'/',strlen(argv[0]));
 353         program_name = p==NULL?argv[0]:(p+1);
 354         //parse arguments
 355         int next_option;
 356         /* A string listing valid short options letters. */
 357         const char* short_options="p:e:P:l:vzh";
 358         /* An array describing valid long options. */
 359         const struct option long_options[] = {
 360                 { "pid", 0, NULL, 'p' },
 361                 { "exe", 1, NULL, 'e' },
 362                 { "path", 0, NULL, 'P' },
 363                 { "limit", 0, NULL, 'l' },
 364                 { "verbose", 0, NULL, 'v' },
 365                 { "lazy", 0, NULL, 'z' },
 366                 { "help", 0, NULL, 'h' },
 367                 { NULL, 0, NULL, 0 }
 368         };
 369         //argument variables
 370         const char *exe=NULL;
 371         const char *path=NULL;
 372         int perclimit=0;
 373         int pid_ok=0;
 374         int process_ok=0;
 375         int limit_ok=0;
 376
 377         do {
 378                 next_option = getopt_long (argc, argv, short_options,long_options, NULL);
 379                 switch(next_option) {
 380                         case 'p':
 381                                 pid=atoi(optarg);
 382                                 pid_ok=1;
 383                                 break;
 384                         case 'e':
 385                                 exe=optarg;
 386                                 process_ok=1;
 387                                 break;
 388                         case 'P':
 389                                 path=optarg;
 390                                 process_ok=1;
 391                                 break;
 392                         case 'l':
 393                                 perclimit=atoi(optarg);
 394                                 limit_ok=1;
 395                                 break;
 396                         case 'v':
 397                                 verbose=1;
 398                                 break;
 399                         case 'z':
 400                                 lazy=1;
 401                                 break;
 402                         case 'h':
 403                                 print_usage (stdout, 1);
 404                                 break;
 405                         case '?':
 406                                 print_usage (stderr, 1);
 407                                 break;
 408                         case -1:
 409                                 break;
 410                         default:
 411                                 abort();
 412                 }
 413         } while(next_option != -1);
 414
 415         if (!process_ok && !pid_ok) {
 416                 fprintf(stderr,"Error: You must specify a target process\n");
 417                 print_usage (stderr, 1);
 418                 exit(1);
 419         }
 420         if ((exe!=NULL && path!=NULL) || (pid_ok && (exe!=NULL || path!=NULL))) {
 421                 fprintf(stderr,"Error: You must specify exactly one target process\n");
 422                 print_usage (stderr, 1);
 423                 exit(1);
 424         }
 425         if (!limit_ok) {
 426                 fprintf(stderr,"Error: You must specify a cpu limit\n");
 427                 print_usage (stderr, 1);
 428                 exit(1);
 429         }
 430         float limit=perclimit/100.0;
 431         if (limit<0 || limit >1) {
 432                 fprintf(stderr,"Error: limit must be in the range 0-100\n");
 433                 print_usage (stderr, 1);
 434                 exit(1);
 435         }
 436         //parameters are all ok!
 437         signal(SIGINT,quit);
 438         signal(SIGTERM,quit);
 439
 440         //time quantum in microseconds. it's splitted in a working period and a sleeping one
 441         int period=100000;
 442         struct timespec twork,tsleep;   //working and sleeping intervals
 443         memset(&twork,0,sizeof(struct timespec));
 444         memset(&tsleep,0,sizeof(struct timespec));
 445
 446 wait_for_process:
 447
 448         //look for the target process..or wait for it
 449         if (exe!=NULL)
 450                 pid=getpidof(exe);
 451         else if (path!=NULL)
 452                 pid=getpidof(path);
 453         else {
 454                 waitforpid(pid);
 455         }
 456         //process detected...let's play
 457
 458         //init compute_cpu_usage internal stuff
 459         compute_cpu_usage(0,0,NULL);
 460         //main loop counter
 461         int i=0;
 462
 463         struct timespec startwork,endwork;
 464         long workingtime=0;             //last working time in microseconds
 465
 466         if (verbose) print_caption();
 467
 468         float pcpu_avg=0;
 469
 470         //here we should already have high priority, for time precision
 471         while(1) {
 472
 473                 //estimate how much the controlled process is using the cpu in its working interval
 474                 struct cpu_usage cu;
 475                 if (compute_cpu_usage(pid,workingtime,&cu)==-1) {
 476                         fprintf(stderr,"Process %d dead!\n",pid);
 477                         if (lazy) exit(2);
 478                         //wait until our process appears
 479                         goto wait_for_process;
 480                 }
 481
 482                 //cpu actual usage of process (range 0-1)
 483                 float pcpu=cu.pcpu;
 484                 //rate at which we are keeping active the process (range 0-1)
 485                 float workingrate=cu.workingrate;
 486
 487                 //adjust work and sleep time slices
 488                 if (pcpu>0) {
 489                         twork.tv_nsec=min(period*limit*1000/pcpu*workingrate,period*1000);
 490                 }
 491                 else if (pcpu==0) {
 492                         twork.tv_nsec=period*1000;
 493                 }
 494                 else if (pcpu==-1) {
 495                         //not yet a valid idea of cpu usage
 496                         pcpu=limit;
 497                         workingrate=limit;
 498                         twork.tv_nsec=min(period*limit*1000,period*1000);
 499                 }
 500                 tsleep.tv_nsec=period*1000-twork.tv_nsec;
 501
 502                 //update average usage
 503                 pcpu_avg=(pcpu_avg*i+pcpu)/(i+1);
 504
 505                 if (verbose && i%10==0 && i>0) {
 506                         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);
 507                 }
 508
 509                 if (limit<1 && limit>0) {
 510                         //resume process
 511                         if (kill(pid,SIGCONT)!=0) {
 512                                 fprintf(stderr,"Process %d dead!\n",pid);
 513                                 if (lazy) exit(2);
 514                                 //wait until our process appears
 515                                 goto wait_for_process;
 516                         }
 517                 }
 518
 519                 clock_gettime(CLOCK_REALTIME,&startwork);
 520                 nanosleep(&twork,NULL);         //now process is working
 521                 clock_gettime(CLOCK_REALTIME,&endwork);
 522                 workingtime=timediff(&endwork,&startwork);
 523
 524                 if (limit<1) {
 525                         //stop process, it has worked enough
 526                         if (kill(pid,SIGSTOP)!=0) {
 527                                 fprintf(stderr,"Process %d dead!\n",pid);
 528                                 if (lazy) exit(2);
 529                                 //wait until our process appears
 530                                 goto wait_for_process;
 531                         }
 532                         nanosleep(&tsleep,NULL);        //now process is sleeping
 533                 }
 534                 i++;
 535         }
 536
 537 }