user-tools/multithrottler

   1 #!/usr/bin/env perl
   2
   3 use lib '/usr/lib/multithrottler';
   4 use Data::Dumper;
   5 use Throttler;
   6 use YAML;
   7 use Getopt::Long;
   8 Getopt::Long::Configure('no_ignore_case');
   9 no if ($] >= 5.018), 'warnings' => 'experimental::smartmatch';
  10
  11 sub warnxf
  12 {
  13         my $_ = shift;
  14         if(defined and $_ ne '')
  15         {
  16                 $_ .= "\n" unless /\n$/;
  17                 s/^/multithrottler: /mg;
  18         }
  19         return $_;
  20 }
  21 sub warnx
  22 {
  23         print STDERR warnxf(@_);
  24 }
  25 sub errx
  26 {
  27         $! = 255;
  28         die warnxf(@_);
  29 }
  30
  31 $OverLimitStatusCode = 11;
  32 $Itemkey = '.';
  33 $Usage = "Usage:
  34   multithrottler [--rate <MAXITEMS>/<INTERVAL> | --reset] --zone <FILE>
  35   multithrottler [-r RATE | -Z] --zone <FILE> [-F] [-c CODE] [--item <KEY>] <COMMAND> [<ARGS>]
  36   multithrottler [--info] [--dump] --zone <FILE>
  37
  38 Options:
  39   -r, --rate      define how many items are permitted within a given interval
  40                   (in seconds). you can just simply define/redefine rate or
  41                   define/redefine rate and run command in a single invocation.
  42                   not need to define rate if you already have the zone file.
  43                   if you redefine the rate, counters get reset.
  44   -n, --no-redefine-rate
  45   --no-redefine   define rate only if the zone does not exists yet,
  46                   i.e. prevent overwrite. use with --rate.
  47   -f, -z, --zonefile
  48   --zone, --file  where to store throttling records
  49   -Z, --reset     reset counters in the zone
  50   -F, --force     force insert item. but throttle is neccessary.
  51   -c, --code      status code for throttling (default is $OverLimitStatusCode)
  52   -i, -k, --itemkey
  53   --item, --key   key of the item to insert (default is \".\")
  54   --info          display rate info
  55   --dump          dump zone records
  56 ";
  57 %ThrottlerParams = ('backend' => 'YAML');
  58
  59
  60 sub readRate
  61 {
  62         eval {
  63                 my $zonedata = YAML::LoadFile($Zonefile);
  64                 my $rate = $zonedata->{'chain'}->{'max_items'} . '/' . $zonedata->{'chain'}->{'interval'};
  65                 parseRate($rate) or die "Can not parse rate info";
  66                 1;
  67         } or errx "Could not load YAML '$Zonefile'. Delete and create again.\n$@";
  68 }
  69 sub parseRate
  70 {
  71         $Rate = $_[0];
  72         if($Rate =~ /^(\d+)\/(\d+)$/) {
  73                 $ThrottlerParams{'max_items'} = $1;
  74                 $ThrottlerParams{'interval'} = $2;
  75         } else {
  76                 return '';
  77         }
  78         return 1;
  79 }
  80 sub cast_throttler
  81 {
  82         $Throttler = Throttler->new(%ThrottlerParams) or errx "Could not cast Throttler instance";
  83 }
  84
  85
  86 GetOptions(
  87         'r|rate=s' => sub{ parseRate($_[1]) or die; },
  88         'n|no-redefine|no-redefine-rate' => \$noRedefineRate,
  89         'Z|reset' => \$ThrottlerParams{'reset'},
  90         'F|force' => \$doForce,
  91         'f|z|file|zone|zonefile=s' => \$Zonefile,
  92         'i|k|item|key|itemkey=s' => \$Itemkey,
  93         'c|code=i' => \$OverLimitStatusCode,
  94         'info' => \$doInfo,
  95         'dump' => \$doDump,
  96 )
  97 or die $Usage;
  98
  99 $ThrottlerParams{'backend_options'}->{'db_file'} = $Zonefile;
 100 $doReset = $ThrottlerParams{'reset'};
 101 ($Command, @Args) = @ARGV;
 102
 103
 104
 105 if(defined $Zonefile)
 106 {
 107         my $mode = '';
 108         $mode .= 'r' if defined $Rate;
 109         $mode .= 'Z' if $doReset;
 110         $mode .= 'C' if defined $Command;
 111         $mode .= 'F' if $doForce;
 112         $mode .= 'x' if $doInfo or $doDump;
 113
 114         if(not($mode ~~ ['r', 'rC', 'rCF', 'Z', 'C', 'CF', 'x', 'ZC', 'ZCF']))
 115         {
 116                 errx "Wrong combination of options. ($mode)";
 117         }
 118
 119         if($mode eq 'x')
 120         {
 121                 readRate;
 122
 123                 if($doInfo)
 124                 {
 125                         print "rate=$Rate\n";
 126                 }
 127                 if($doDump)
 128                 {
 129                         cast_throttler();
 130                         print $Throttler->buckets_dump;
 131                 }
 132                 exit 0;
 133         }
 134
 135         if($noRedefineRate)
 136         {
 137                 eval { readRate; 1; }
 138         }
 139
 140         if(!defined $Rate)
 141         {
 142                 readRate;
 143         }
 144
 145         cast_throttler();
 146
 147         if($mode =~ /C/)
 148         {
 149                 my $allowed;
 150                 for my $attempt (1, 2)
 151                 {
 152                         if(eval {
 153                                 if($Throttler->try_push('key' => $Itemkey, 'force' => $doForce))
 154                                 {
 155                                         $allowed = 1;
 156                                 }
 157                                 else
 158                                 {
 159                                         $allowed = 0;
 160                                 }
 161                                 1;
 162                         })
 163                         {
 164                                 last;
 165                         }
 166                         else
 167                         {
 168                                 warnx "Throttler module error\n$@";
 169                                 if($attempt == 1)
 170                                 {
 171                                         warnx "Resetting Throttler zone file";
 172                                         # if there was Throttler module exception, try to reset the backend repairing it so.
 173                                         undef $Throttler;
 174                                         $ThrottlerParams{'reset'} = 1;
 175                                         cast_throttler();
 176                                 }
 177                                 else
 178                                 {
 179                                         errx;
 180                                 }
 181                         }
 182                 }
 183
 184                 if($allowed)
 185                 {
 186                         exec {$Command} $Command, @Args;
 187                         my ($errno, $errstr) = (int $!, $!);
 188                         warn "$0: $Command: $errstr\n";
 189                         exit 125+$errno;
 190                 }
 191                 else
 192                 {
 193                         exit $OverLimitStatusCode;
 194                 }
 195         }
 196         else
 197         {
 198                 # it was either just a rate setup or counter reset
 199                 exit 0;
 200         }
 201 }
 202
 203 die $Usage;
 204
 205
 206 __END__
 207
 208 =pod
 209
 210 =head1 NAME
 211
 212 multithrottler - Run given command if not reached the defined rate limit
 213
 214 =cut
 215