1 #include <string.h>
2
3 #include "events.h"
4 #include "monoclock.h"
5
6 #include "tsnetwork.h"
7 #include "tsnetwork_internal.h"
8
9 /* Bandwidth limiting state. */
10 struct bwlimit {
11 double Bps;
12 double bucket;
13 void * timer_cookie;
14 int suspended;
15 };
16 static struct bwlimit limit_read;
17 static struct bwlimit limit_write;
18
19 /* Last time tokens were added to buckets. */
20 static struct timeval tlast;
21 static int tlast_set = 0;
22
23 /* Initialized? */
24 static int initdone = 0;
25
26 /* Update the internal state. */
27 static int poke(void);
28
29 /* Initialize if necessary. */
30 static void
init(void)31 init(void)
32 {
33
34 /* Limit to 1 GBps by default. */
35 limit_read.Bps = 1000000000.0;
36
37 /* 2s burst. */
38 limit_read.bucket = 2 * limit_read.Bps;
39
40 /* No timer yet. */
41 limit_read.timer_cookie = NULL;
42
43 /* Traffic not suspended. */
44 limit_read.suspended = 0;
45
46 /* Write state is the same as the read state. */
47 memcpy(&limit_write, &limit_read, sizeof(struct bwlimit));
48
49 /* We've been initialized! */
50 initdone = 1;
51 }
52
53 /* Timer wakeup. */
54 static int
callback_timer(void * cookie)55 callback_timer(void * cookie)
56 {
57 struct bwlimit * l = cookie;
58
59 /* This timer is no longer running. */
60 l->timer_cookie = NULL;
61
62 /* Update state. */
63 return (poke());
64 }
65
66 /* Update the state for one direction. */
67 static int
pokeone(struct bwlimit * l,double t,int op)68 pokeone(struct bwlimit * l, double t, int op)
69 {
70 double waketime;
71
72 /* Add tokens to the bucket. */
73 l->bucket += l->Bps * t;
74
75 /* Overflow the bucket at 2 seconds of bandwidth. */
76 if (l->bucket > 2 * l->Bps)
77 l->bucket = 2 * l->Bps;
78
79 /* Do we need to re-enable traffic? */
80 if ((l->bucket >= 1460) && (l->suspended != 0)) {
81 /* Allow traffic to pass. */
82 if (network_register_resume(op))
83 goto err0;
84 l->suspended = 0;
85 }
86
87 /* Do we need to block traffic? */
88 if ((l->bucket < 1460) && (l->suspended == 0)) {
89 /* Stop traffic from running. */
90 if (network_register_suspend(op))
91 goto err0;
92 l->suspended = 1;
93 }
94
95 /* If traffic is running, we don't need a timer. */
96 if ((l->suspended == 0) && (l->timer_cookie != NULL)) {
97 events_timer_cancel(l->timer_cookie);
98 l->timer_cookie = NULL;
99 }
100
101 /* If traffic is suspended, we need a timer. */
102 if ((l->suspended == 1) && (l->timer_cookie == NULL)) {
103 /* Wait 10 ms or for 1460 bytes of quota. */
104 waketime = (1460 - l->bucket) / l->Bps;
105 if (waketime < 0.01)
106 waketime = 0.01;
107
108 /* Register a timer. */
109 if ((l->timer_cookie =
110 events_timer_register_double(callback_timer,
111 l, waketime)) == NULL)
112 goto err0;
113 }
114
115 /* Success! */
116 return (0);
117
118 err0:
119 /* Failure! */
120 return (-1);
121 }
122
123 /* Update the internal state. */
124 static int
poke(void)125 poke(void)
126 {
127 struct timeval tnow;
128 double t;
129
130 /* Get the current time. */
131 if (monoclock_get(&tnow))
132 goto err0;
133
134 /* Compute the duration since the last poke. */
135 if (tlast_set)
136 t = (tnow.tv_sec - tlast.tv_sec) +
137 (tnow.tv_usec - tlast.tv_usec) * 0.000001;
138 else
139 t = 0.0;
140
141 /* Poke each direction. */
142 if (pokeone(&limit_read, t, NETWORK_OP_READ))
143 goto err0;
144 if (pokeone(&limit_write, t, NETWORK_OP_WRITE))
145 goto err0;
146
147 /* We have been poked. */
148 memcpy(&tlast, &tnow, sizeof(struct timeval));
149 tlast_set = 1;
150
151 /* Success! */
152 return (0);
153
154 err0:
155 /* Failure! */
156 return (-1);
157 }
158
159 /**
160 * network_bwlimit(down, up):
161 * Set the bandwidth rate limit to ${down} bytes per second of read bandwidth
162 * and ${up} bytes per second of write bandwidth. The values ${down} and
163 * ${up} must be between 8000 and 10^9.
164 */
165 void
network_bwlimit(double down,double up)166 network_bwlimit(double down, double up)
167 {
168
169 /* Initialize if necessary. */
170 if (!initdone)
171 init();
172
173 /* Record these values for future reference. */
174 limit_read.Bps = down;
175 limit_write.Bps = up;
176
177 /* Don't allow more than a 2 s burst. */
178 if (limit_read.bucket > 2 * limit_read.Bps)
179 limit_read.bucket = 2 * limit_read.Bps;
180 if (limit_write.bucket > 2 * limit_write.Bps)
181 limit_write.bucket = 2 * limit_write.Bps;
182 }
183
184 /**
185 * network_bwlimit_get(op, len):
186 * Get the amount of instantaneously allowed bandwidth for ${op} operations.
187 */
188 int
network_bwlimit_get(int op,size_t * len)189 network_bwlimit_get(int op, size_t * len)
190 {
191
192 /* Initialize if necessary. */
193 if (!initdone)
194 init();
195
196 /* Update state. */
197 if (poke())
198 goto err0;
199
200 /* Return the appropriate value. */
201 if (op == NETWORK_OP_READ)
202 *len = (size_t)limit_read.bucket;
203 else
204 *len = (size_t)limit_write.bucket;
205
206 /*
207 * If the allowed bandwidth is less than one normal-sized TCP segment,
208 * force it to zero, in order to avoid silly windowing.
209 */
210 if (*len < 1460)
211 *len = 0;
212
213 /* Success! */
214 return (0);
215
216 err0:
217 /* Failure! */
218 return (-1);
219 }
220
221 /**
222 * network_bwlimit_eat(op, len):
223 * Consume ${len} bytes of bandwidth quota for ${op} operations.
224 */
225 int
network_bwlimit_eat(int op,size_t len)226 network_bwlimit_eat(int op, size_t len)
227 {
228
229 /* Initialize if necessary. */
230 if (!initdone)
231 init();
232
233 /* Eat tokens from the bucket. */
234 if (op == NETWORK_OP_READ)
235 limit_read.bucket -= len;
236 else
237 limit_write.bucket -= len;
238
239 /* Update state. */
240 return (poke());
241 }
242