qLibc
qtokenbucket.c
Go to the documentation of this file.
1 /******************************************************************************
2  * qLibc
3  *
4  * Copyright (c) 2010-2015 Seungyoung Kim.
5  * All rights reserved.
6  *
7  * Redistribution and use in source and binary forms, with or without
8  * modification, are permitted provided that the following conditions are met:
9  *
10  * 1. Redistributions of source code must retain the above copyright notice,
11  * this list of conditions and the following disclaimer.
12  * 2. Redistributions in binary form must reproduce the above copyright notice,
13  * this list of conditions and the following disclaimer in the documentation
14  * and/or other materials provided with the distribution.
15  *
16  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
17  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
19  * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
20  * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
21  * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
22  * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
23  * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
24  * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
25  * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
26  * POSSIBILITY OF SUCH DAMAGE.
27  *****************************************************************************/
28 
29 /**
30  * @file qtokenbucket.c Token Bucket implementation.
31  *
32  * Current implementation is not thread-safe.
33  *
34  * More information about token-bucket:
35  * http://en.wikipedia.org/wiki/Token_bucket
36  *
37  * @code
38  * qtokenbucket_t bucket;
39  * qtokenbucket_init(&bucket, 500, 1000, 1000);
40  * while (1) {
41  * if (qtokenbucket_consume(&bucket, 1) == false) {
42  * // Bucket is empty. Let's wait
43  * usleep(qtokenbucket_waittime(&bucket, 1) * 1000);
44  * continue;
45  * }
46  * // Got a token. Let's do something here.
47  * do_something();
48  * }
49  * @endcode
50  */
51 
52 #include "extensions/qtokenbucket.h"
53 
54 #include <stdio.h>
55 #include <stdbool.h>
56 #include <string.h>
57 #include <sys/time.h>
58 #include "utilities/qtime.h"
59 #include "qinternal.h"
60 
61 #ifndef _DOXYGEN_SKIP
62 /* internal functions */
63 static void refill_tokens(qtokenbucket_t *bucket);
64 #endif
65 
66 /**
67  * Initialize the token bucket.
68  *
69  * @param init_tokens
70  * the initial number of tokens.
71  * @param max_tokens
72  * maximum number of tokens in the bucket.
73  * @param tokens_per_sec
74  * number of tokens to fill per a second.
75  */
76 void qtokenbucket_init(qtokenbucket_t *bucket, int init_tokens, int max_tokens,
77  int tokens_per_sec) {
78  memset(bucket, 0, sizeof(qtokenbucket_t));
79  bucket->tokens = init_tokens;
80  bucket->max_tokens = max_tokens;
81  bucket->tokens_per_sec = tokens_per_sec;
82  bucket->last_fill = qtime_current_milli();
83 }
84 
85 /**
86  * Consume tokens from the bucket.
87  *
88  * @param bucket tockenbucket object.
89  * @param tokens number of tokens to request.
90  *
91  * @return return true if there are enough tokens, otherwise false.
92  */
93 bool qtokenbucket_consume(qtokenbucket_t *bucket, int tokens) {
94  refill_tokens(bucket);
95  if (bucket->tokens < tokens) {
96  return false;
97  }
98  bucket->tokens -= tokens;
99  return true;
100 }
101 
102 /**
103  * Get the estimate time until given number of token is ready.
104  *
105  * @param tokens number of tokens
106  *
107  * @return estimated milliseconds
108  */
109 long qtokenbucket_waittime(qtokenbucket_t *bucket, int tokens) {
110  refill_tokens(bucket);
111  if (bucket->tokens >= tokens) {
112  return 0;
113  }
114  int tokens_needed = tokens - (int)bucket->tokens;
115  double estimate_milli = (1000 * tokens_needed) / bucket->tokens_per_sec;
116  estimate_milli += ((1000 * tokens_needed) % bucket->tokens_per_sec) ? 1 : 0;
117  return estimate_milli;
118 }
119 
120 #ifndef _DOXYGEN_SKIP
121 /**
122  * Refill tokens.
123  */
124 static void refill_tokens(qtokenbucket_t *bucket) {
125  long now = qtime_current_milli();
126  if (bucket->tokens < bucket->max_tokens) {
127  double new_tokens = (now - bucket->last_fill) * 0.001
128  * bucket->tokens_per_sec;
129  bucket->tokens =
130  ((bucket->tokens + new_tokens) < bucket->max_tokens) ?
131  (bucket->tokens + new_tokens) : bucket->max_tokens;
132  }
133  bucket->last_fill = now;
134 }
135 #endif // _DOXYGEN_SKIP
long qtokenbucket_waittime(qtokenbucket_t *bucket, int tokens)
Get the estimate time until given number of token is ready.
Definition: qtokenbucket.c:109
bool qtokenbucket_consume(qtokenbucket_t *bucket, int tokens)
Consume tokens from the bucket.
Definition: qtokenbucket.c:93
long qtime_current_milli(void)
Returns the current time in milliseconds.
Definition: qtime.c:49
void qtokenbucket_init(qtokenbucket_t *bucket, int init_tokens, int max_tokens, int tokens_per_sec)
Initialize the token bucket.
Definition: qtokenbucket.c:76