1 /* $NetBSD: ratelimiter.c,v 1.7 2022/09/23 12:15:33 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 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 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 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 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 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 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 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 244 isc_ratelimiter_shutdown(isc_ratelimiter_t *rl) { 245 isc_event_t *ev; 246 isc_task_t *task; 247 248 REQUIRE(rl != NULL); 249 250 LOCK(&rl->lock); 251 rl->state = isc_ratelimiter_shuttingdown; 252 (void)isc_timer_reset(rl->timer, isc_timertype_inactive, NULL, NULL, 253 false); 254 while ((ev = ISC_LIST_HEAD(rl->pending)) != NULL) { 255 task = ev->ev_sender; 256 ISC_LIST_UNLINK(rl->pending, ev, ev_ratelink); 257 ev->ev_attributes |= ISC_EVENTATTR_CANCELED; 258 isc_task_send(task, &ev); 259 } 260 task = NULL; 261 isc_task_attach(rl->task, &task); 262 isc_timer_detach(&rl->timer); 263 264 /* 265 * Send an event to our task. The delivery of this event 266 * indicates that no more timer events will be delivered. 267 */ 268 ev = &rl->shutdownevent; 269 isc_task_send(rl->task, &ev); 270 271 UNLOCK(&rl->lock); 272 } 273 274 static void 275 ratelimiter_shutdowncomplete(isc_task_t *task, isc_event_t *event) { 276 isc_ratelimiter_t *rl = (isc_ratelimiter_t *)event->ev_arg; 277 278 UNUSED(task); 279 280 isc_ratelimiter_detach(&rl); 281 isc_task_detach(&task); 282 } 283 284 static void 285 ratelimiter_free(isc_ratelimiter_t *rl) { 286 isc_refcount_destroy(&rl->references); 287 isc_mutex_destroy(&rl->lock); 288 isc_mem_put(rl->mctx, rl, sizeof(*rl)); 289 } 290 291 void 292 isc_ratelimiter_attach(isc_ratelimiter_t *source, isc_ratelimiter_t **target) { 293 REQUIRE(source != NULL); 294 REQUIRE(target != NULL && *target == NULL); 295 296 isc_refcount_increment(&source->references); 297 298 *target = source; 299 } 300 301 void 302 isc_ratelimiter_detach(isc_ratelimiter_t **rlp) { 303 isc_ratelimiter_t *rl; 304 305 REQUIRE(rlp != NULL && *rlp != NULL); 306 307 rl = *rlp; 308 *rlp = NULL; 309 310 if (isc_refcount_decrement(&rl->references) == 1) { 311 ratelimiter_free(rl); 312 } 313 } 314 315 isc_result_t 316 isc_ratelimiter_stall(isc_ratelimiter_t *rl) { 317 isc_result_t result = ISC_R_SUCCESS; 318 319 REQUIRE(rl != NULL); 320 321 LOCK(&rl->lock); 322 switch (rl->state) { 323 case isc_ratelimiter_shuttingdown: 324 result = ISC_R_SHUTTINGDOWN; 325 break; 326 case isc_ratelimiter_ratelimited: 327 result = isc_timer_reset(rl->timer, isc_timertype_inactive, 328 NULL, NULL, false); 329 RUNTIME_CHECK(result == ISC_R_SUCCESS); 330 FALLTHROUGH; 331 case isc_ratelimiter_idle: 332 case isc_ratelimiter_stalled: 333 rl->state = isc_ratelimiter_stalled; 334 break; 335 } 336 UNLOCK(&rl->lock); 337 return (result); 338 } 339 340 isc_result_t 341 isc_ratelimiter_release(isc_ratelimiter_t *rl) { 342 isc_result_t result = ISC_R_SUCCESS; 343 344 REQUIRE(rl != NULL); 345 346 LOCK(&rl->lock); 347 switch (rl->state) { 348 case isc_ratelimiter_shuttingdown: 349 result = ISC_R_SHUTTINGDOWN; 350 break; 351 case isc_ratelimiter_stalled: 352 if (!ISC_LIST_EMPTY(rl->pending)) { 353 result = isc_timer_reset(rl->timer, 354 isc_timertype_ticker, NULL, 355 &rl->interval, false); 356 if (result == ISC_R_SUCCESS) { 357 rl->state = isc_ratelimiter_ratelimited; 358 } 359 } else { 360 rl->state = isc_ratelimiter_idle; 361 } 362 break; 363 case isc_ratelimiter_ratelimited: 364 case isc_ratelimiter_idle: 365 break; 366 } 367 UNLOCK(&rl->lock); 368 return (result); 369 } 370