ucw/tbf.c

   1 /*
   2  *      UCW Library -- Rate Limiting based on the Token Bucket Filter
   3  *
   4  *      (c) 2009 Martin Mares <mj@ucw.cz>
   5  *
   6  *      This software may be freely distributed and used according to the terms
   7  *      of the GNU Lesser General Public License.
   8  */
   9
  10 #include <ucw/lib.h>
  11 #include <ucw/tbf.h>
  12
  13 void
  14 tbf_init(struct token_bucket_filter *f)
  15 {
  16   if (!f->burst)
  17     f->burst = MAX(2*f->rate, 1);
  18   f->last_hit = 0;
  19   f->bucket = f->burst;
  20 }
  21
  22 int
  23 tbf_limit(struct token_bucket_filter *f, timestamp_t now)
  24 {
  25   timestamp_t delta_t = now - f->last_hit;
  26   f->last_hit = now;
  27
  28   double b = f->bucket + f->rate * delta_t / 1000;
  29   b = MIN(b, f->burst);
  30   if (b >= 1)
  31     {
  32       uint dropped = f->drop_count;
  33       f->bucket = b - 1;
  34       f->drop_count = 0;
  35       return dropped;
  36     }
  37   else
  38     {
  39       f->bucket = b;
  40       f->drop_count++;
  41       return -f->drop_count;
  42     }
  43 }
  44
  45 #ifdef TEST
  46
  47 int main(void)
  48 {
  49   struct token_bucket_filter t = { .rate = 1, .burst = 2 };
  50   tbf_init(&t);
  51   for (timestamp_t now = 0; now < 3000; now += 77)
  52     {
  53       int res = tbf_limit(&t, now);
  54       msg(L_DEBUG, "t=%u result=%d bucket=%f", (uint) now, res, t.bucket);
  55     }
  56   return 0;
  57 }
  58
  59 #endif