1 /*
2 * Copyright (C) Internet Systems Consortium, Inc. ("ISC")
3 *
4 * SPDX-License-Identifier: MPL-2.0
5 *
6 * This Source Code Form is subject to the terms of the Mozilla Public
7 * License, v. 2.0. If a copy of the MPL was not distributed with this
8 * file, you can obtain one at https://mozilla.org/MPL/2.0/.
9 *
10 * See the COPYRIGHT file distributed with this work for additional
11 * information regarding copyright ownership.
12 */
13
14 /*! \file */
15
16 #include <inttypes.h>
17 #include <stdbool.h>
18
19 #include <isc/mem.h>
20 #include <isc/ratelimiter.h>
21 #include <isc/refcount.h>
22 #include <isc/task.h>
23 #include <isc/time.h>
24 #include <isc/timer.h>
25 #include <isc/util.h>
26
27 typedef enum {
28 isc_ratelimiter_stalled = 0,
29 isc_ratelimiter_ratelimited = 1,
30 isc_ratelimiter_idle = 2,
31 isc_ratelimiter_shuttingdown = 3
32 } isc_ratelimiter_state_t;
33
34 struct isc_ratelimiter {
35 isc_mem_t *mctx;
36 isc_mutex_t lock;
37 isc_refcount_t references;
38 isc_task_t *task;
39 isc_timer_t *timer;
40 isc_interval_t interval;
41 uint32_t pertic;
42 bool pushpop;
43 isc_ratelimiter_state_t state;
44 isc_event_t shutdownevent;
45 ISC_LIST(isc_event_t) pending;
46 };
47
48 #define ISC_RATELIMITEREVENT_SHUTDOWN (ISC_EVENTCLASS_RATELIMITER + 1)
49
50 static void
51 ratelimiter_tick(isc_task_t *task, isc_event_t *event);
52
53 static void
54 ratelimiter_shutdowncomplete(isc_task_t *task, isc_event_t *event);
55
56 isc_result_t
isc_ratelimiter_create(isc_mem_t * mctx,isc_timermgr_t * timermgr,isc_task_t * task,isc_ratelimiter_t ** ratelimiterp)57 isc_ratelimiter_create(isc_mem_t *mctx, isc_timermgr_t *timermgr,
58 isc_task_t *task, isc_ratelimiter_t **ratelimiterp) {
59 isc_result_t result;
60 isc_ratelimiter_t *rl;
61 INSIST(ratelimiterp != NULL && *ratelimiterp == NULL);
62
63 rl = isc_mem_get(mctx, sizeof(*rl));
64 *rl = (isc_ratelimiter_t){
65 .mctx = mctx,
66 .task = task,
67 .pertic = 1,
68 .state = isc_ratelimiter_idle,
69 };
70
71 isc_refcount_init(&rl->references, 1);
72 isc_interval_set(&rl->interval, 0, 0);
73 ISC_LIST_INIT(rl->pending);
74
75 isc_mutex_init(&rl->lock);
76
77 result = isc_timer_create(timermgr, isc_timertype_inactive, NULL, NULL,
78 rl->task, ratelimiter_tick, rl, &rl->timer);
79 if (result != ISC_R_SUCCESS) {
80 goto free_mutex;
81 }
82
83 /*
84 * Increment the reference count to indicate that we may
85 * (soon) have events outstanding.
86 */
87 isc_refcount_increment(&rl->references);
88
89 ISC_EVENT_INIT(&rl->shutdownevent, sizeof(isc_event_t), 0, NULL,
90 ISC_RATELIMITEREVENT_SHUTDOWN,
91 ratelimiter_shutdowncomplete, rl, rl, NULL, NULL);
92
93 *ratelimiterp = rl;
94 return (ISC_R_SUCCESS);
95
96 free_mutex:
97 isc_refcount_decrementz(&rl->references);
98 isc_refcount_destroy(&rl->references);
99 isc_mutex_destroy(&rl->lock);
100 isc_mem_put(mctx, rl, sizeof(*rl));
101 return (result);
102 }
103
104 isc_result_t
isc_ratelimiter_setinterval(isc_ratelimiter_t * rl,isc_interval_t * interval)105 isc_ratelimiter_setinterval(isc_ratelimiter_t *rl, isc_interval_t *interval) {
106 isc_result_t result = ISC_R_SUCCESS;
107
108 REQUIRE(rl != NULL);
109 REQUIRE(interval != NULL);
110
111 LOCK(&rl->lock);
112 rl->interval = *interval;
113 /*
114 * If the timer is currently running, change its rate.
115 */
116 if (rl->state == isc_ratelimiter_ratelimited) {
117 result = isc_timer_reset(rl->timer, isc_timertype_ticker, NULL,
118 &rl->interval, false);
119 }
120 UNLOCK(&rl->lock);
121 return (result);
122 }
123
124 void
isc_ratelimiter_setpertic(isc_ratelimiter_t * rl,uint32_t pertic)125 isc_ratelimiter_setpertic(isc_ratelimiter_t *rl, uint32_t pertic) {
126 REQUIRE(rl != NULL);
127
128 if (pertic == 0) {
129 pertic = 1;
130 }
131 rl->pertic = pertic;
132 }
133
134 void
isc_ratelimiter_setpushpop(isc_ratelimiter_t * rl,bool pushpop)135 isc_ratelimiter_setpushpop(isc_ratelimiter_t *rl, bool pushpop) {
136 REQUIRE(rl != NULL);
137
138 rl->pushpop = pushpop;
139 }
140
141 isc_result_t
isc_ratelimiter_enqueue(isc_ratelimiter_t * rl,isc_task_t * task,isc_event_t ** eventp)142 isc_ratelimiter_enqueue(isc_ratelimiter_t *rl, isc_task_t *task,
143 isc_event_t **eventp) {
144 isc_result_t result = ISC_R_SUCCESS;
145 isc_event_t *ev;
146
147 REQUIRE(rl != NULL);
148 REQUIRE(task != NULL);
149 REQUIRE(eventp != NULL && *eventp != NULL);
150 ev = *eventp;
151 REQUIRE(ev->ev_sender == NULL);
152
153 LOCK(&rl->lock);
154 if (rl->state == isc_ratelimiter_ratelimited ||
155 rl->state == isc_ratelimiter_stalled)
156 {
157 ev->ev_sender = task;
158 *eventp = NULL;
159 if (rl->pushpop) {
160 ISC_LIST_PREPEND(rl->pending, ev, ev_ratelink);
161 } else {
162 ISC_LIST_APPEND(rl->pending, ev, ev_ratelink);
163 }
164 } else if (rl->state == isc_ratelimiter_idle) {
165 result = isc_timer_reset(rl->timer, isc_timertype_ticker, NULL,
166 &rl->interval, false);
167 if (result == ISC_R_SUCCESS) {
168 ev->ev_sender = task;
169 rl->state = isc_ratelimiter_ratelimited;
170 }
171 } else {
172 INSIST(rl->state == isc_ratelimiter_shuttingdown);
173 result = ISC_R_SHUTTINGDOWN;
174 }
175 UNLOCK(&rl->lock);
176 if (*eventp != NULL && result == ISC_R_SUCCESS) {
177 isc_task_send(task, eventp);
178 }
179 return (result);
180 }
181
182 isc_result_t
isc_ratelimiter_dequeue(isc_ratelimiter_t * rl,isc_event_t * event)183 isc_ratelimiter_dequeue(isc_ratelimiter_t *rl, isc_event_t *event) {
184 isc_result_t result = ISC_R_SUCCESS;
185
186 REQUIRE(rl != NULL);
187 REQUIRE(event != NULL);
188
189 LOCK(&rl->lock);
190 if (ISC_LINK_LINKED(event, ev_ratelink)) {
191 ISC_LIST_UNLINK(rl->pending, event, ev_ratelink);
192 event->ev_sender = NULL;
193 } else {
194 result = ISC_R_NOTFOUND;
195 }
196 UNLOCK(&rl->lock);
197 return (result);
198 }
199
200 static void
ratelimiter_tick(isc_task_t * task,isc_event_t * event)201 ratelimiter_tick(isc_task_t *task, isc_event_t *event) {
202 isc_ratelimiter_t *rl = (isc_ratelimiter_t *)event->ev_arg;
203 isc_event_t *p;
204 uint32_t pertic;
205
206 UNUSED(task);
207
208 isc_event_free(&event);
209
210 pertic = rl->pertic;
211 while (pertic != 0) {
212 pertic--;
213 LOCK(&rl->lock);
214 p = ISC_LIST_HEAD(rl->pending);
215 if (p != NULL) {
216 /*
217 * There is work to do. Let's do it after unlocking.
218 */
219 ISC_LIST_UNLINK(rl->pending, p, ev_ratelink);
220 } else {
221 /*
222 * No work left to do. Stop the timer so that we don't
223 * waste resources by having it fire periodically.
224 */
225 isc_result_t result = isc_timer_reset(
226 rl->timer, isc_timertype_inactive, NULL, NULL,
227 false);
228 RUNTIME_CHECK(result == ISC_R_SUCCESS);
229 rl->state = isc_ratelimiter_idle;
230 pertic = 0; /* Force the loop to exit. */
231 }
232 UNLOCK(&rl->lock);
233 if (p != NULL) {
234 isc_task_t *evtask = p->ev_sender;
235 isc_task_send(evtask, &p);
236 }
237 INSIST(p == NULL);
238 }
239 }
240
241 void
isc_ratelimiter_shutdown(isc_ratelimiter_t * rl)242 isc_ratelimiter_shutdown(isc_ratelimiter_t *rl) {
243 isc_event_t *ev;
244 isc_task_t *task;
245
246 REQUIRE(rl != NULL);
247
248 LOCK(&rl->lock);
249 rl->state = isc_ratelimiter_shuttingdown;
250 (void)isc_timer_reset(rl->timer, isc_timertype_inactive, NULL, NULL,
251 false);
252 while ((ev = ISC_LIST_HEAD(rl->pending)) != NULL) {
253 task = ev->ev_sender;
254 ISC_LIST_UNLINK(rl->pending, ev, ev_ratelink);
255 ev->ev_attributes |= ISC_EVENTATTR_CANCELED;
256 isc_task_send(task, &ev);
257 }
258 task = NULL;
259 isc_task_attach(rl->task, &task);
260 isc_timer_detach(&rl->timer);
261
262 /*
263 * Send an event to our task. The delivery of this event
264 * indicates that no more timer events will be delivered.
265 */
266 ev = &rl->shutdownevent;
267 isc_task_send(rl->task, &ev);
268
269 UNLOCK(&rl->lock);
270 }
271
272 static void
ratelimiter_shutdowncomplete(isc_task_t * task,isc_event_t * event)273 ratelimiter_shutdowncomplete(isc_task_t *task, isc_event_t *event) {
274 isc_ratelimiter_t *rl = (isc_ratelimiter_t *)event->ev_arg;
275
276 UNUSED(task);
277
278 isc_ratelimiter_detach(&rl);
279 isc_task_detach(&task);
280 }
281
282 static void
ratelimiter_free(isc_ratelimiter_t * rl)283 ratelimiter_free(isc_ratelimiter_t *rl) {
284 isc_refcount_destroy(&rl->references);
285 isc_mutex_destroy(&rl->lock);
286 isc_mem_put(rl->mctx, rl, sizeof(*rl));
287 }
288
289 void
isc_ratelimiter_attach(isc_ratelimiter_t * source,isc_ratelimiter_t ** target)290 isc_ratelimiter_attach(isc_ratelimiter_t *source, isc_ratelimiter_t **target) {
291 REQUIRE(source != NULL);
292 REQUIRE(target != NULL && *target == NULL);
293
294 isc_refcount_increment(&source->references);
295
296 *target = source;
297 }
298
299 void
isc_ratelimiter_detach(isc_ratelimiter_t ** rlp)300 isc_ratelimiter_detach(isc_ratelimiter_t **rlp) {
301 isc_ratelimiter_t *rl;
302
303 REQUIRE(rlp != NULL && *rlp != NULL);
304
305 rl = *rlp;
306 *rlp = NULL;
307
308 if (isc_refcount_decrement(&rl->references) == 1) {
309 ratelimiter_free(rl);
310 }
311 }
312
313 isc_result_t
isc_ratelimiter_stall(isc_ratelimiter_t * rl)314 isc_ratelimiter_stall(isc_ratelimiter_t *rl) {
315 isc_result_t result = ISC_R_SUCCESS;
316
317 REQUIRE(rl != NULL);
318
319 LOCK(&rl->lock);
320 switch (rl->state) {
321 case isc_ratelimiter_shuttingdown:
322 result = ISC_R_SHUTTINGDOWN;
323 break;
324 case isc_ratelimiter_ratelimited:
325 result = isc_timer_reset(rl->timer, isc_timertype_inactive,
326 NULL, NULL, false);
327 RUNTIME_CHECK(result == ISC_R_SUCCESS);
328 /* FALLTHROUGH */
329 case isc_ratelimiter_idle:
330 case isc_ratelimiter_stalled:
331 rl->state = isc_ratelimiter_stalled;
332 break;
333 }
334 UNLOCK(&rl->lock);
335 return (result);
336 }
337
338 isc_result_t
isc_ratelimiter_release(isc_ratelimiter_t * rl)339 isc_ratelimiter_release(isc_ratelimiter_t *rl) {
340 isc_result_t result = ISC_R_SUCCESS;
341
342 REQUIRE(rl != NULL);
343
344 LOCK(&rl->lock);
345 switch (rl->state) {
346 case isc_ratelimiter_shuttingdown:
347 result = ISC_R_SHUTTINGDOWN;
348 break;
349 case isc_ratelimiter_stalled:
350 if (!ISC_LIST_EMPTY(rl->pending)) {
351 result = isc_timer_reset(rl->timer,
352 isc_timertype_ticker, NULL,
353 &rl->interval, false);
354 if (result == ISC_R_SUCCESS) {
355 rl->state = isc_ratelimiter_ratelimited;
356 }
357 } else {
358 rl->state = isc_ratelimiter_idle;
359 }
360 break;
361 case isc_ratelimiter_ratelimited:
362 case isc_ratelimiter_idle:
363 break;
364 }
365 UNLOCK(&rl->lock);
366 return (result);
367 }
368