1258945Sroberto/*
2258945Sroberto * Copyright (C) 2004, 2005, 2007  Internet Systems Consortium, Inc. ("ISC")
3258945Sroberto * Copyright (C) 1999-2002  Internet Software Consortium.
4258945Sroberto *
5258945Sroberto * Permission to use, copy, modify, and/or distribute this software for any
6258945Sroberto * purpose with or without fee is hereby granted, provided that the above
7258945Sroberto * copyright notice and this permission notice appear in all copies.
8258945Sroberto *
9258945Sroberto * THE SOFTWARE IS PROVIDED "AS IS" AND ISC DISCLAIMS ALL WARRANTIES WITH
10258945Sroberto * REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
11258945Sroberto * AND FITNESS.  IN NO EVENT SHALL ISC BE LIABLE FOR ANY SPECIAL, DIRECT,
12258945Sroberto * INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
13258945Sroberto * LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
14258945Sroberto * OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
15258945Sroberto * PERFORMANCE OF THIS SOFTWARE.
16258945Sroberto */
17258945Sroberto
18258945Sroberto/* $Id: ratelimiter.c,v 1.25 2007/06/19 23:47:17 tbox Exp $ */
19258945Sroberto
20258945Sroberto/*! \file */
21258945Sroberto
22258945Sroberto#include <config.h>
23258945Sroberto
24258945Sroberto#include <isc/mem.h>
25258945Sroberto#include <isc/ratelimiter.h>
26258945Sroberto#include <isc/task.h>
27258945Sroberto#include <isc/time.h>
28258945Sroberto#include <isc/timer.h>
29258945Sroberto#include <isc/util.h>
30258945Sroberto
31258945Srobertotypedef enum {
32258945Sroberto	isc_ratelimiter_stalled = 0,
33258945Sroberto	isc_ratelimiter_ratelimited = 1,
34258945Sroberto	isc_ratelimiter_idle = 2,
35258945Sroberto	isc_ratelimiter_shuttingdown = 3
36258945Sroberto} isc_ratelimiter_state_t;
37258945Sroberto
38258945Srobertostruct isc_ratelimiter {
39258945Sroberto	isc_mem_t *		mctx;
40258945Sroberto	isc_mutex_t		lock;
41258945Sroberto	int			refs;
42258945Sroberto	isc_task_t *		task;
43258945Sroberto	isc_timer_t *		timer;
44258945Sroberto	isc_interval_t		interval;
45258945Sroberto	isc_uint32_t		pertic;
46258945Sroberto	isc_ratelimiter_state_t	state;
47258945Sroberto	isc_event_t		shutdownevent;
48258945Sroberto	ISC_LIST(isc_event_t)	pending;
49258945Sroberto};
50258945Sroberto
51258945Sroberto#define ISC_RATELIMITEREVENT_SHUTDOWN (ISC_EVENTCLASS_RATELIMITER + 1)
52258945Sroberto
53258945Srobertostatic void
54258945Srobertoratelimiter_tick(isc_task_t *task, isc_event_t *event);
55258945Sroberto
56258945Srobertostatic void
57258945Srobertoratelimiter_shutdowncomplete(isc_task_t *task, isc_event_t *event);
58258945Sroberto
59258945Srobertoisc_result_t
60258945Srobertoisc_ratelimiter_create(isc_mem_t *mctx, isc_timermgr_t *timermgr,
61258945Sroberto		       isc_task_t *task, isc_ratelimiter_t **ratelimiterp)
62258945Sroberto{
63258945Sroberto	isc_result_t result;
64258945Sroberto	isc_ratelimiter_t *rl;
65258945Sroberto	INSIST(ratelimiterp != NULL && *ratelimiterp == NULL);
66258945Sroberto
67258945Sroberto	rl = isc_mem_get(mctx, sizeof(*rl));
68258945Sroberto	if (rl == NULL)
69258945Sroberto		return ISC_R_NOMEMORY;
70258945Sroberto	rl->mctx = mctx;
71258945Sroberto	rl->refs = 1;
72258945Sroberto	rl->task = task;
73258945Sroberto	isc_interval_set(&rl->interval, 0, 0);
74258945Sroberto	rl->timer = NULL;
75258945Sroberto	rl->pertic = 1;
76258945Sroberto	rl->state = isc_ratelimiter_idle;
77258945Sroberto	ISC_LIST_INIT(rl->pending);
78258945Sroberto
79258945Sroberto	result = isc_mutex_init(&rl->lock);
80258945Sroberto	if (result != ISC_R_SUCCESS)
81258945Sroberto		goto free_mem;
82258945Sroberto	result = isc_timer_create(timermgr, isc_timertype_inactive,
83258945Sroberto				  NULL, NULL, rl->task, ratelimiter_tick,
84258945Sroberto				  rl, &rl->timer);
85258945Sroberto	if (result != ISC_R_SUCCESS)
86258945Sroberto		goto free_mutex;
87258945Sroberto
88258945Sroberto	/*
89258945Sroberto	 * Increment the reference count to indicate that we may
90258945Sroberto	 * (soon) have events outstanding.
91258945Sroberto	 */
92258945Sroberto	rl->refs++;
93258945Sroberto
94258945Sroberto	ISC_EVENT_INIT(&rl->shutdownevent,
95258945Sroberto		       sizeof(isc_event_t),
96258945Sroberto		       0, NULL, ISC_RATELIMITEREVENT_SHUTDOWN,
97258945Sroberto		       ratelimiter_shutdowncomplete, rl, rl, NULL, NULL);
98258945Sroberto
99258945Sroberto	*ratelimiterp = rl;
100258945Sroberto	return (ISC_R_SUCCESS);
101258945Sroberto
102258945Srobertofree_mutex:
103258945Sroberto	DESTROYLOCK(&rl->lock);
104258945Srobertofree_mem:
105258945Sroberto	isc_mem_put(mctx, rl, sizeof(*rl));
106258945Sroberto	return (result);
107258945Sroberto}
108258945Sroberto
109258945Srobertoisc_result_t
110258945Srobertoisc_ratelimiter_setinterval(isc_ratelimiter_t *rl, isc_interval_t *interval) {
111258945Sroberto	isc_result_t result = ISC_R_SUCCESS;
112258945Sroberto	LOCK(&rl->lock);
113258945Sroberto	rl->interval = *interval;
114258945Sroberto	/*
115258945Sroberto	 * If the timer is currently running, change its rate.
116258945Sroberto	 */
117258945Sroberto        if (rl->state == isc_ratelimiter_ratelimited) {
118258945Sroberto		result = isc_timer_reset(rl->timer, isc_timertype_ticker, NULL,
119258945Sroberto					 &rl->interval, ISC_FALSE);
120258945Sroberto	}
121258945Sroberto	UNLOCK(&rl->lock);
122258945Sroberto	return (result);
123258945Sroberto}
124258945Sroberto
125258945Srobertovoid
126258945Srobertoisc_ratelimiter_setpertic(isc_ratelimiter_t *rl, isc_uint32_t pertic) {
127258945Sroberto	if (pertic == 0)
128258945Sroberto		pertic = 1;
129258945Sroberto	rl->pertic = pertic;
130258945Sroberto}
131258945Sroberto
132258945Srobertoisc_result_t
133258945Srobertoisc_ratelimiter_enqueue(isc_ratelimiter_t *rl, isc_task_t *task,
134258945Sroberto			isc_event_t **eventp)
135258945Sroberto{
136258945Sroberto	isc_result_t result = ISC_R_SUCCESS;
137258945Sroberto	isc_event_t *ev;
138258945Sroberto
139258945Sroberto	REQUIRE(eventp != NULL && *eventp != NULL);
140258945Sroberto	REQUIRE(task != NULL);
141258945Sroberto	ev = *eventp;
142258945Sroberto	REQUIRE(ev->ev_sender == NULL);
143258945Sroberto
144258945Sroberto	LOCK(&rl->lock);
145258945Sroberto        if (rl->state == isc_ratelimiter_ratelimited ||
146258945Sroberto	    rl->state == isc_ratelimiter_stalled) {
147258945Sroberto		isc_event_t *ev = *eventp;
148258945Sroberto		ev->ev_sender = task;
149258945Sroberto                ISC_LIST_APPEND(rl->pending, ev, ev_link);
150258945Sroberto		*eventp = NULL;
151258945Sroberto        } else if (rl->state == isc_ratelimiter_idle) {
152258945Sroberto		result = isc_timer_reset(rl->timer, isc_timertype_ticker, NULL,
153258945Sroberto					 &rl->interval, ISC_FALSE);
154258945Sroberto		if (result == ISC_R_SUCCESS) {
155258945Sroberto			ev->ev_sender = task;
156258945Sroberto			rl->state = isc_ratelimiter_ratelimited;
157258945Sroberto		}
158258945Sroberto	} else {
159258945Sroberto		INSIST(rl->state == isc_ratelimiter_shuttingdown);
160258945Sroberto		result = ISC_R_SHUTTINGDOWN;
161258945Sroberto	}
162258945Sroberto	UNLOCK(&rl->lock);
163258945Sroberto	if (*eventp != NULL && result == ISC_R_SUCCESS)
164258945Sroberto		isc_task_send(task, eventp);
165258945Sroberto	return (result);
166258945Sroberto}
167258945Sroberto
168258945Srobertostatic void
169258945Srobertoratelimiter_tick(isc_task_t *task, isc_event_t *event) {
170258945Sroberto	isc_result_t result = ISC_R_SUCCESS;
171258945Sroberto	isc_ratelimiter_t *rl = (isc_ratelimiter_t *)event->ev_arg;
172258945Sroberto	isc_event_t *p;
173258945Sroberto	isc_uint32_t pertic;
174258945Sroberto
175258945Sroberto	UNUSED(task);
176258945Sroberto
177258945Sroberto	isc_event_free(&event);
178258945Sroberto
179258945Sroberto	pertic = rl->pertic;
180258945Sroberto        while (pertic != 0) {
181258945Sroberto		pertic--;
182258945Sroberto		LOCK(&rl->lock);
183258945Sroberto		p = ISC_LIST_HEAD(rl->pending);
184258945Sroberto		if (p != NULL) {
185258945Sroberto			/*
186258945Sroberto			 * There is work to do.  Let's do it after unlocking.
187258945Sroberto			 */
188258945Sroberto			ISC_LIST_UNLINK(rl->pending, p, ev_link);
189258945Sroberto		} else {
190258945Sroberto			/*
191258945Sroberto			 * No work left to do.  Stop the timer so that we don't
192258945Sroberto			 * waste resources by having it fire periodically.
193258945Sroberto			 */
194258945Sroberto			result = isc_timer_reset(rl->timer,
195258945Sroberto						 isc_timertype_inactive,
196258945Sroberto						 NULL, NULL, ISC_FALSE);
197258945Sroberto			RUNTIME_CHECK(result == ISC_R_SUCCESS);
198258945Sroberto			rl->state = isc_ratelimiter_idle;
199258945Sroberto			pertic = 0;	/* Force the loop to exit. */
200258945Sroberto		}
201258945Sroberto		UNLOCK(&rl->lock);
202258945Sroberto		if (p != NULL) {
203258945Sroberto			isc_task_t *evtask = p->ev_sender;
204258945Sroberto			isc_task_send(evtask, &p);
205258945Sroberto		}
206258945Sroberto		INSIST(p == NULL);
207258945Sroberto	}
208258945Sroberto}
209258945Sroberto
210258945Srobertovoid
211258945Srobertoisc_ratelimiter_shutdown(isc_ratelimiter_t *rl) {
212258945Sroberto	isc_event_t *ev;
213258945Sroberto	isc_task_t *task;
214258945Sroberto	LOCK(&rl->lock);
215258945Sroberto	rl->state = isc_ratelimiter_shuttingdown;
216258945Sroberto	(void)isc_timer_reset(rl->timer, isc_timertype_inactive,
217258945Sroberto			      NULL, NULL, ISC_FALSE);
218258945Sroberto	while ((ev = ISC_LIST_HEAD(rl->pending)) != NULL) {
219258945Sroberto		ISC_LIST_UNLINK(rl->pending, ev, ev_link);
220258945Sroberto		ev->ev_attributes |= ISC_EVENTATTR_CANCELED;
221258945Sroberto		task = ev->ev_sender;
222258945Sroberto		isc_task_send(task, &ev);
223258945Sroberto	}
224258945Sroberto	isc_timer_detach(&rl->timer);
225258945Sroberto	/*
226258945Sroberto	 * Send an event to our task.  The delivery of this event
227258945Sroberto	 * indicates that no more timer events will be delivered.
228258945Sroberto	 */
229258945Sroberto	ev = &rl->shutdownevent;
230258945Sroberto	isc_task_send(rl->task, &ev);
231258945Sroberto
232258945Sroberto	UNLOCK(&rl->lock);
233258945Sroberto}
234258945Sroberto
235258945Srobertostatic void
236258945Srobertoratelimiter_shutdowncomplete(isc_task_t *task, isc_event_t *event) {
237258945Sroberto	isc_ratelimiter_t *rl = (isc_ratelimiter_t *)event->ev_arg;
238258945Sroberto
239258945Sroberto	UNUSED(task);
240258945Sroberto
241258945Sroberto	isc_ratelimiter_detach(&rl);
242258945Sroberto}
243258945Sroberto
244258945Srobertostatic void
245258945Srobertoratelimiter_free(isc_ratelimiter_t *rl) {
246258945Sroberto	DESTROYLOCK(&rl->lock);
247258945Sroberto	isc_mem_put(rl->mctx, rl, sizeof(*rl));
248258945Sroberto}
249258945Sroberto
250258945Srobertovoid
251258945Srobertoisc_ratelimiter_attach(isc_ratelimiter_t *source, isc_ratelimiter_t **target) {
252258945Sroberto	REQUIRE(source != NULL);
253258945Sroberto	REQUIRE(target != NULL && *target == NULL);
254258945Sroberto
255258945Sroberto	LOCK(&source->lock);
256258945Sroberto	REQUIRE(source->refs > 0);
257258945Sroberto	source->refs++;
258258945Sroberto	INSIST(source->refs > 0);
259258945Sroberto	UNLOCK(&source->lock);
260258945Sroberto	*target = source;
261258945Sroberto}
262258945Sroberto
263258945Srobertovoid
264258945Srobertoisc_ratelimiter_detach(isc_ratelimiter_t **rlp) {
265258945Sroberto	isc_ratelimiter_t *rl = *rlp;
266258945Sroberto	isc_boolean_t free_now = ISC_FALSE;
267258945Sroberto
268258945Sroberto	LOCK(&rl->lock);
269258945Sroberto	REQUIRE(rl->refs > 0);
270258945Sroberto	rl->refs--;
271258945Sroberto	if (rl->refs == 0)
272258945Sroberto		free_now = ISC_TRUE;
273258945Sroberto	UNLOCK(&rl->lock);
274258945Sroberto
275258945Sroberto	if (free_now)
276258945Sroberto		ratelimiter_free(rl);
277258945Sroberto
278258945Sroberto	*rlp = NULL;
279258945Sroberto}
280258945Sroberto
281258945Srobertoisc_result_t
282258945Srobertoisc_ratelimiter_stall(isc_ratelimiter_t *rl) {
283258945Sroberto	isc_result_t result = ISC_R_SUCCESS;
284258945Sroberto
285258945Sroberto	LOCK(&rl->lock);
286258945Sroberto	switch (rl->state) {
287258945Sroberto	case isc_ratelimiter_shuttingdown:
288258945Sroberto		result = ISC_R_SHUTTINGDOWN;
289258945Sroberto		break;
290258945Sroberto	case isc_ratelimiter_ratelimited:
291258945Sroberto		result = isc_timer_reset(rl->timer, isc_timertype_inactive,
292258945Sroberto				 	 NULL, NULL, ISC_FALSE);
293258945Sroberto		RUNTIME_CHECK(result == ISC_R_SUCCESS);
294258945Sroberto	case isc_ratelimiter_idle:
295258945Sroberto	case isc_ratelimiter_stalled:
296258945Sroberto		rl->state = isc_ratelimiter_stalled;
297258945Sroberto		break;
298258945Sroberto	}
299258945Sroberto	UNLOCK(&rl->lock);
300258945Sroberto	return (result);
301258945Sroberto}
302258945Sroberto
303258945Srobertoisc_result_t
304258945Srobertoisc_ratelimiter_release(isc_ratelimiter_t *rl) {
305258945Sroberto	isc_result_t result = ISC_R_SUCCESS;
306258945Sroberto
307258945Sroberto	LOCK(&rl->lock);
308258945Sroberto	switch (rl->state) {
309258945Sroberto	case isc_ratelimiter_shuttingdown:
310258945Sroberto		result = ISC_R_SHUTTINGDOWN;
311258945Sroberto		break;
312258945Sroberto	case isc_ratelimiter_stalled:
313258945Sroberto		if (!ISC_LIST_EMPTY(rl->pending)) {
314258945Sroberto			result = isc_timer_reset(rl->timer,
315258945Sroberto						 isc_timertype_ticker, NULL,
316258945Sroberto						 &rl->interval, ISC_FALSE);
317258945Sroberto			if (result == ISC_R_SUCCESS)
318258945Sroberto				rl->state = isc_ratelimiter_ratelimited;
319258945Sroberto		} else
320258945Sroberto			rl->state = isc_ratelimiter_idle;
321258945Sroberto		break;
322258945Sroberto	case isc_ratelimiter_ratelimited:
323258945Sroberto	case isc_ratelimiter_idle:
324258945Sroberto		break;
325258945Sroberto	}
326258945Sroberto	UNLOCK(&rl->lock);
327258945Sroberto	return (result);
328258945Sroberto}
329