1 /* $NetBSD: ratelimiter.c,v 1.5 2020/05/24 19:46:26 christos Exp $ */ 2 3 /* 4 * Copyright (C) Internet Systems Consortium, Inc. ("ISC") 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 http://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 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_mutex_destroy(&rl->lock); 98 isc_mem_put(mctx, rl, sizeof(*rl)); 99 return (result); 100 } 101 102 isc_result_t 103 isc_ratelimiter_setinterval(isc_ratelimiter_t *rl, isc_interval_t *interval) { 104 isc_result_t result = ISC_R_SUCCESS; 105 106 REQUIRE(rl != NULL); 107 REQUIRE(interval != NULL); 108 109 LOCK(&rl->lock); 110 rl->interval = *interval; 111 /* 112 * If the timer is currently running, change its rate. 113 */ 114 if (rl->state == isc_ratelimiter_ratelimited) { 115 result = isc_timer_reset(rl->timer, isc_timertype_ticker, NULL, 116 &rl->interval, false); 117 } 118 UNLOCK(&rl->lock); 119 return (result); 120 } 121 122 void 123 isc_ratelimiter_setpertic(isc_ratelimiter_t *rl, uint32_t pertic) { 124 REQUIRE(rl != NULL); 125 126 if (pertic == 0) { 127 pertic = 1; 128 } 129 rl->pertic = pertic; 130 } 131 132 void 133 isc_ratelimiter_setpushpop(isc_ratelimiter_t *rl, bool pushpop) { 134 REQUIRE(rl != NULL); 135 136 rl->pushpop = pushpop; 137 } 138 139 isc_result_t 140 isc_ratelimiter_enqueue(isc_ratelimiter_t *rl, isc_task_t *task, 141 isc_event_t **eventp) { 142 isc_result_t result = ISC_R_SUCCESS; 143 isc_event_t *ev; 144 145 REQUIRE(rl != NULL); 146 REQUIRE(task != NULL); 147 REQUIRE(eventp != NULL && *eventp != NULL); 148 ev = *eventp; 149 REQUIRE(ev->ev_sender == NULL); 150 151 LOCK(&rl->lock); 152 if (rl->state == isc_ratelimiter_ratelimited || 153 rl->state == isc_ratelimiter_stalled) 154 { 155 ev->ev_sender = task; 156 *eventp = NULL; 157 if (rl->pushpop) { 158 ISC_LIST_PREPEND(rl->pending, ev, ev_ratelink); 159 } else { 160 ISC_LIST_APPEND(rl->pending, ev, ev_ratelink); 161 } 162 } else if (rl->state == isc_ratelimiter_idle) { 163 result = isc_timer_reset(rl->timer, isc_timertype_ticker, NULL, 164 &rl->interval, false); 165 if (result == ISC_R_SUCCESS) { 166 ev->ev_sender = task; 167 rl->state = isc_ratelimiter_ratelimited; 168 } 169 } else { 170 INSIST(rl->state == isc_ratelimiter_shuttingdown); 171 result = ISC_R_SHUTTINGDOWN; 172 } 173 UNLOCK(&rl->lock); 174 if (*eventp != NULL && result == ISC_R_SUCCESS) { 175 isc_task_send(task, eventp); 176 } 177 return (result); 178 } 179 180 isc_result_t 181 isc_ratelimiter_dequeue(isc_ratelimiter_t *rl, isc_event_t *event) { 182 isc_result_t result = ISC_R_SUCCESS; 183 184 REQUIRE(rl != NULL); 185 REQUIRE(event != NULL); 186 187 LOCK(&rl->lock); 188 if (ISC_LINK_LINKED(event, ev_ratelink)) { 189 ISC_LIST_UNLINK(rl->pending, event, ev_ratelink); 190 event->ev_sender = NULL; 191 } else { 192 result = ISC_R_NOTFOUND; 193 } 194 UNLOCK(&rl->lock); 195 return (result); 196 } 197 198 static void 199 ratelimiter_tick(isc_task_t *task, isc_event_t *event) { 200 isc_ratelimiter_t *rl = (isc_ratelimiter_t *)event->ev_arg; 201 isc_event_t *p; 202 uint32_t pertic; 203 204 UNUSED(task); 205 206 isc_event_free(&event); 207 208 pertic = rl->pertic; 209 while (pertic != 0) { 210 pertic--; 211 LOCK(&rl->lock); 212 p = ISC_LIST_HEAD(rl->pending); 213 if (p != NULL) { 214 /* 215 * There is work to do. Let's do it after unlocking. 216 */ 217 ISC_LIST_UNLINK(rl->pending, p, ev_ratelink); 218 } else { 219 /* 220 * No work left to do. Stop the timer so that we don't 221 * waste resources by having it fire periodically. 222 */ 223 isc_result_t result = isc_timer_reset( 224 rl->timer, isc_timertype_inactive, NULL, NULL, 225 false); 226 RUNTIME_CHECK(result == ISC_R_SUCCESS); 227 rl->state = isc_ratelimiter_idle; 228 pertic = 0; /* Force the loop to exit. */ 229 } 230 UNLOCK(&rl->lock); 231 if (p != NULL) { 232 isc_task_t *evtask = p->ev_sender; 233 isc_task_send(evtask, &p); 234 } 235 INSIST(p == NULL); 236 } 237 } 238 239 void 240 isc_ratelimiter_shutdown(isc_ratelimiter_t *rl) { 241 isc_event_t *ev; 242 243 REQUIRE(rl != NULL); 244 245 LOCK(&rl->lock); 246 rl->state = isc_ratelimiter_shuttingdown; 247 (void)isc_timer_reset(rl->timer, isc_timertype_inactive, NULL, NULL, 248 false); 249 while ((ev = ISC_LIST_HEAD(rl->pending)) != NULL) { 250 isc_task_t *task = ev->ev_sender; 251 ISC_LIST_UNLINK(rl->pending, ev, ev_ratelink); 252 ev->ev_attributes |= ISC_EVENTATTR_CANCELED; 253 isc_task_send(task, &ev); 254 } 255 isc_timer_detach(&rl->timer); 256 257 /* 258 * Send an event to our task. The delivery of this event 259 * indicates that no more timer events will be delivered. 260 */ 261 ev = &rl->shutdownevent; 262 isc_task_send(rl->task, &ev); 263 264 UNLOCK(&rl->lock); 265 } 266 267 static void 268 ratelimiter_shutdowncomplete(isc_task_t *task, isc_event_t *event) { 269 isc_ratelimiter_t *rl = (isc_ratelimiter_t *)event->ev_arg; 270 271 UNUSED(task); 272 273 isc_ratelimiter_detach(&rl); 274 } 275 276 static void 277 ratelimiter_free(isc_ratelimiter_t *rl) { 278 isc_mutex_destroy(&rl->lock); 279 isc_mem_put(rl->mctx, rl, sizeof(*rl)); 280 } 281 282 void 283 isc_ratelimiter_attach(isc_ratelimiter_t *source, isc_ratelimiter_t **target) { 284 REQUIRE(source != NULL); 285 REQUIRE(target != NULL && *target == NULL); 286 287 isc_refcount_increment(&source->references); 288 289 *target = source; 290 } 291 292 void 293 isc_ratelimiter_detach(isc_ratelimiter_t **rlp) { 294 isc_ratelimiter_t *rl; 295 296 REQUIRE(rlp != NULL && *rlp != NULL); 297 298 rl = *rlp; 299 *rlp = NULL; 300 301 if (isc_refcount_decrement(&rl->references) == 1) { 302 ratelimiter_free(rl); 303 } 304 } 305 306 isc_result_t 307 isc_ratelimiter_stall(isc_ratelimiter_t *rl) { 308 isc_result_t result = ISC_R_SUCCESS; 309 310 REQUIRE(rl != NULL); 311 312 LOCK(&rl->lock); 313 switch (rl->state) { 314 case isc_ratelimiter_shuttingdown: 315 result = ISC_R_SHUTTINGDOWN; 316 break; 317 case isc_ratelimiter_ratelimited: 318 result = isc_timer_reset(rl->timer, isc_timertype_inactive, 319 NULL, NULL, false); 320 RUNTIME_CHECK(result == ISC_R_SUCCESS); 321 /* FALLTHROUGH */ 322 case isc_ratelimiter_idle: 323 case isc_ratelimiter_stalled: 324 rl->state = isc_ratelimiter_stalled; 325 break; 326 } 327 UNLOCK(&rl->lock); 328 return (result); 329 } 330 331 isc_result_t 332 isc_ratelimiter_release(isc_ratelimiter_t *rl) { 333 isc_result_t result = ISC_R_SUCCESS; 334 335 REQUIRE(rl != NULL); 336 337 LOCK(&rl->lock); 338 switch (rl->state) { 339 case isc_ratelimiter_shuttingdown: 340 result = ISC_R_SHUTTINGDOWN; 341 break; 342 case isc_ratelimiter_stalled: 343 if (!ISC_LIST_EMPTY(rl->pending)) { 344 result = isc_timer_reset(rl->timer, 345 isc_timertype_ticker, NULL, 346 &rl->interval, false); 347 if (result == ISC_R_SUCCESS) { 348 rl->state = isc_ratelimiter_ratelimited; 349 } 350 } else { 351 rl->state = isc_ratelimiter_idle; 352 } 353 break; 354 case isc_ratelimiter_ratelimited: 355 case isc_ratelimiter_idle: 356 break; 357 } 358 UNLOCK(&rl->lock); 359 return (result); 360 } 361