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