lib/token-bucket.c

   1 /*
   2  * Copyright (c) 2012 Nicira, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at:
   7  *
   8  *     http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include <config.h>
  18
  19 #include "token-bucket.h"
  20
  21 #include "poll-loop.h"
  22 #include "sat-math.h"
  23 #include "timeval.h"
  24 #include "util.h"
  25
  26 /* Initializes 'tb' to accumulate 'rate' tokens per millisecond, with a
  27  * maximum of 'burst' tokens.
  28  *
  29  * The token bucket is initially full.
  30  *
  31  * It may be more convenient to use TOKEN_BUCKET_INIT. */
  32 void
  33 token_bucket_init(struct token_bucket *tb,
  34                   unsigned int rate, unsigned int burst)
  35 {
  36     tb->rate = rate;
  37     tb->burst = burst;
  38     tb->tokens = 0;
  39     tb->last_fill = LLONG_MIN;
  40 }
  41
  42 /* Changes 'tb' to accumulate 'rate' tokens per millisecond, with a maximum of
  43  * 'burst' tokens.
  44  *
  45  * 'tb' must already have been initialized with TOKEN_BUCKET_INIT or
  46  * token_bucket_init(). */
  47 void
  48 token_bucket_set(struct token_bucket *tb,
  49                  unsigned int rate, unsigned int burst)
  50 {
  51     tb->rate = rate;
  52     tb->burst = burst;
  53     if (burst > tb->tokens) {
  54         tb->tokens = burst;
  55     }
  56 }
  57
  58 /* Attempts to remove 'n' tokens from 'tb'.  Returns true if successful, false
  59  * if 'tb' contained fewer than 'n' tokens (and thus 'n' tokens could not be
  60  * removed) . */
  61 bool
  62 token_bucket_withdraw(struct token_bucket *tb, unsigned int n)
  63 {
  64     if (tb->tokens < n) {
  65         long long int now = time_msec();
  66         if (now > tb->last_fill) {
  67             unsigned long long int elapsed_ull
  68                 = (unsigned long long int) now - tb->last_fill;
  69             unsigned int elapsed = MIN(UINT_MAX, elapsed_ull);
  70             unsigned int add = sat_mul(tb->rate, elapsed);
  71             unsigned int tokens = sat_add(tb->tokens, add);
  72             tb->tokens = MIN(tokens, tb->burst);
  73             tb->last_fill = now;
  74         }
  75
  76         if (tb->tokens < n) {
  77             return false;
  78         }
  79     }
  80
  81     tb->tokens -= n;
  82     return true;
  83 }
  84
  85 /* Causes the poll loop to wake up when at least 'n' tokens will be available
  86  * for withdrawal from 'tb'. */
  87 void
  88 token_bucket_wait(struct token_bucket *tb, unsigned int n)
  89 {
  90     if (tb->tokens >= n) {
  91         poll_immediate_wake();
  92     } else {
  93         unsigned int need = n - tb->tokens;
  94         poll_timer_wait_until(tb->last_fill + need / tb->rate + 1);
  95     }
  96 }