1 /* $NetBSD: ratelimiter.c,v 1.10 2025/01/26 16:25:38 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/async.h> 22 #include <isc/loop.h> 23 #include <isc/magic.h> 24 #include <isc/mem.h> 25 #include <isc/ratelimiter.h> 26 #include <isc/refcount.h> 27 #include <isc/time.h> 28 #include <isc/timer.h> 29 #include <isc/util.h> 30 31 typedef enum { 32 isc_ratelimiter_ratelimited = 0, 33 isc_ratelimiter_idle = 1, 34 isc_ratelimiter_shuttingdown = 2 35 } isc_ratelimiter_state_t; 36 37 #define RATELIMITER_MAGIC ISC_MAGIC('R', 't', 'L', 'm') 38 #define VALID_RATELIMITER(rl) ISC_MAGIC_VALID(rl, RATELIMITER_MAGIC) 39 40 struct isc_ratelimiter { 41 int magic; 42 isc_mem_t *mctx; 43 isc_loop_t *loop; 44 isc_refcount_t references; 45 isc_mutex_t lock; 46 isc_timer_t *timer; 47 isc_interval_t interval; 48 uint32_t pertic; 49 bool pushpop; 50 isc_ratelimiter_state_t state; 51 ISC_LIST(isc_rlevent_t) pending; 52 }; 53 54 static void 55 isc__ratelimiter_tick(void *arg); 56 57 static void 58 isc__ratelimiter_start(void *arg); 59 60 static void 61 isc__ratelimiter_doshutdown(void *arg); 62 63 void 64 isc_ratelimiter_create(isc_loop_t *loop, isc_ratelimiter_t **rlp) { 65 isc_ratelimiter_t *rl = NULL; 66 isc_mem_t *mctx; 67 68 REQUIRE(loop != NULL); 69 REQUIRE(rlp != NULL && *rlp == NULL); 70 71 mctx = isc_loop_getmctx(loop); 72 73 rl = isc_mem_get(mctx, sizeof(*rl)); 74 *rl = (isc_ratelimiter_t){ 75 .pertic = 1, 76 .state = isc_ratelimiter_idle, 77 .magic = RATELIMITER_MAGIC, 78 }; 79 80 isc_mem_attach(mctx, &rl->mctx); 81 isc_loop_attach(loop, &rl->loop); 82 isc_refcount_init(&rl->references, 1); 83 isc_interval_set(&rl->interval, 0, 0); 84 ISC_LIST_INIT(rl->pending); 85 86 isc_timer_create(rl->loop, isc__ratelimiter_tick, rl, &rl->timer); 87 88 isc_mutex_init(&rl->lock); 89 90 *rlp = rl; 91 } 92 93 void 94 isc_ratelimiter_setinterval(isc_ratelimiter_t *restrict rl, 95 const isc_interval_t *const interval) { 96 REQUIRE(VALID_RATELIMITER(rl)); 97 REQUIRE(interval != NULL); 98 99 LOCK(&rl->lock); 100 rl->interval = *interval; 101 /* The interval will be adjusted on the next tick */ 102 UNLOCK(&rl->lock); 103 } 104 105 void 106 isc_ratelimiter_setpertic(isc_ratelimiter_t *restrict rl, 107 const uint32_t pertic) { 108 REQUIRE(VALID_RATELIMITER(rl)); 109 REQUIRE(pertic > 0); 110 111 LOCK(&rl->lock); 112 rl->pertic = pertic; 113 UNLOCK(&rl->lock); 114 } 115 116 void 117 isc_ratelimiter_setpushpop(isc_ratelimiter_t *restrict rl, const bool pushpop) { 118 REQUIRE(VALID_RATELIMITER(rl)); 119 120 LOCK(&rl->lock); 121 rl->pushpop = pushpop; 122 UNLOCK(&rl->lock); 123 } 124 125 static void 126 isc__ratelimiter_start(void *arg) { 127 isc_ratelimiter_t *rl = arg; 128 isc_interval_t interval; 129 130 REQUIRE(VALID_RATELIMITER(rl)); 131 132 LOCK(&rl->lock); 133 switch (rl->state) { 134 case isc_ratelimiter_ratelimited: 135 /* The first tick happens immediately */ 136 isc_interval_set(&interval, 0, 0); 137 isc_timer_start(rl->timer, isc_timertype_once, &interval); 138 break; 139 case isc_ratelimiter_shuttingdown: 140 /* The ratelimiter is shutting down */ 141 break; 142 case isc_ratelimiter_idle: 143 /* 144 * This could happen if we are changing the interval on the 145 * ratelimiter, but all the events were processed and the timer 146 * was stopped before the new interval could be applied. 147 */ 148 break; 149 default: 150 UNREACHABLE(); 151 } 152 UNLOCK(&rl->lock); 153 isc_ratelimiter_detach(&rl); 154 } 155 156 isc_result_t 157 isc_ratelimiter_enqueue(isc_ratelimiter_t *restrict rl, 158 isc_loop_t *restrict loop, isc_job_cb cb, void *arg, 159 isc_rlevent_t **rlep) { 160 isc_result_t result = ISC_R_SUCCESS; 161 isc_rlevent_t *rle = NULL; 162 163 REQUIRE(VALID_RATELIMITER(rl)); 164 REQUIRE(loop != NULL); 165 REQUIRE(rlep != NULL && *rlep == NULL); 166 167 LOCK(&rl->lock); 168 switch (rl->state) { 169 case isc_ratelimiter_shuttingdown: 170 result = ISC_R_SHUTTINGDOWN; 171 break; 172 case isc_ratelimiter_idle: 173 /* Start the ratelimiter */ 174 isc_ratelimiter_ref(rl); 175 isc_async_run(rl->loop, isc__ratelimiter_start, rl); 176 rl->state = isc_ratelimiter_ratelimited; 177 FALLTHROUGH; 178 case isc_ratelimiter_ratelimited: 179 rle = isc_mem_get(isc_loop_getmctx(loop), sizeof(*rle)); 180 *rle = (isc_rlevent_t){ 181 .cb = cb, 182 .arg = arg, 183 .link = ISC_LINK_INITIALIZER, 184 }; 185 isc_loop_attach(loop, &rle->loop); 186 isc_ratelimiter_attach(rl, &rle->rl); 187 188 if (rl->pushpop) { 189 ISC_LIST_PREPEND(rl->pending, rle, link); 190 } else { 191 ISC_LIST_APPEND(rl->pending, rle, link); 192 } 193 *rlep = rle; 194 break; 195 default: 196 UNREACHABLE(); 197 } 198 UNLOCK(&rl->lock); 199 return result; 200 } 201 202 isc_result_t 203 isc_ratelimiter_dequeue(isc_ratelimiter_t *restrict rl, isc_rlevent_t **rlep) { 204 isc_result_t result = ISC_R_SUCCESS; 205 206 REQUIRE(rl != NULL); 207 REQUIRE(rlep != NULL); 208 209 LOCK(&rl->lock); 210 if (ISC_LINK_LINKED(*rlep, link)) { 211 ISC_LIST_UNLINK(rl->pending, *rlep, link); 212 isc_rlevent_free(rlep); 213 } else { 214 result = ISC_R_NOTFOUND; 215 } 216 UNLOCK(&rl->lock); 217 return result; 218 } 219 220 static void 221 isc__ratelimiter_tick(void *arg) { 222 isc_ratelimiter_t *rl = (isc_ratelimiter_t *)arg; 223 isc_rlevent_t *rle = NULL; 224 uint32_t pertic; 225 ISC_LIST(isc_rlevent_t) pending; 226 227 REQUIRE(VALID_RATELIMITER(rl)); 228 229 ISC_LIST_INIT(pending); 230 231 LOCK(&rl->lock); 232 233 REQUIRE(rl->timer != NULL); 234 235 if (rl->state == isc_ratelimiter_shuttingdown) { 236 INSIST(EMPTY(rl->pending)); 237 goto unlock; 238 } 239 240 pertic = rl->pertic; 241 while (pertic != 0) { 242 rle = ISC_LIST_HEAD(rl->pending); 243 if (rle != NULL) { 244 /* There is work to do. Let's do it after unlocking. */ 245 ISC_LIST_UNLINK(rl->pending, rle, link); 246 ISC_LIST_APPEND(pending, rle, link); 247 } else { 248 /* 249 * We processed all the scheduled work, but there's a 250 * room for at least one more event (we haven't consumed 251 * all of the "pertick"), so we can stop the ratelimiter 252 * now, and don't worry about isc_ratelimiter_enqueue() 253 * sending an extra event immediately. 254 */ 255 rl->state = isc_ratelimiter_idle; 256 break; 257 } 258 pertic--; 259 } 260 261 if (rl->state != isc_ratelimiter_idle) { 262 /* Reschedule the timer */ 263 isc_timer_start(rl->timer, isc_timertype_once, &rl->interval); 264 } 265 unlock: 266 UNLOCK(&rl->lock); 267 268 while ((rle = ISC_LIST_HEAD(pending)) != NULL) { 269 ISC_LIST_UNLINK(pending, rle, link); 270 isc_async_run(rle->loop, rle->cb, rle->arg); 271 } 272 } 273 274 void 275 isc__ratelimiter_doshutdown(void *arg) { 276 isc_ratelimiter_t *rl = arg; 277 278 REQUIRE(VALID_RATELIMITER(rl)); 279 280 LOCK(&rl->lock); 281 INSIST(rl->state == isc_ratelimiter_shuttingdown); 282 INSIST(EMPTY(rl->pending)); 283 284 isc_timer_stop(rl->timer); 285 isc_timer_destroy(&rl->timer); 286 isc_loop_detach(&rl->loop); 287 UNLOCK(&rl->lock); 288 isc_ratelimiter_detach(&rl); 289 } 290 291 void 292 isc_ratelimiter_shutdown(isc_ratelimiter_t *restrict rl) { 293 isc_rlevent_t *rle = NULL; 294 ISC_LIST(isc_rlevent_t) pending; 295 296 REQUIRE(VALID_RATELIMITER(rl)); 297 298 ISC_LIST_INIT(pending); 299 300 LOCK(&rl->lock); 301 if (rl->state != isc_ratelimiter_shuttingdown) { 302 rl->state = isc_ratelimiter_shuttingdown; 303 ISC_LIST_MOVE(pending, rl->pending); 304 isc_ratelimiter_ref(rl); 305 isc_async_run(rl->loop, isc__ratelimiter_doshutdown, rl); 306 } 307 UNLOCK(&rl->lock); 308 309 while ((rle = ISC_LIST_HEAD(pending)) != NULL) { 310 ISC_LIST_UNLINK(pending, rle, link); 311 rle->canceled = true; 312 isc_async_run(rl->loop, rle->cb, rle->arg); 313 } 314 } 315 316 static void 317 ratelimiter_destroy(isc_ratelimiter_t *restrict rl) { 318 LOCK(&rl->lock); 319 REQUIRE(rl->state == isc_ratelimiter_shuttingdown); 320 UNLOCK(&rl->lock); 321 322 isc_mutex_destroy(&rl->lock); 323 isc_mem_putanddetach(&rl->mctx, rl, sizeof(*rl)); 324 } 325 326 void 327 isc_rlevent_free(isc_rlevent_t **rlep) { 328 REQUIRE(rlep != NULL && *rlep != NULL); 329 330 isc_rlevent_t *rle = *rlep; 331 isc_mem_t *mctx = isc_loop_getmctx(rle->loop); 332 333 *rlep = NULL; 334 335 isc_loop_detach(&rle->loop); 336 isc_ratelimiter_detach(&rle->rl); 337 isc_mem_put(mctx, rle, sizeof(*rle)); 338 } 339 340 ISC_REFCOUNT_IMPL(isc_ratelimiter, ratelimiter_destroy); 341