1/*	$NetBSD: ratelimiter.c,v 1.2.6.1 2012/06/05 21:15:07 bouyer Exp $	*/
2
3/*
4 * Copyright (C) 2004, 2005, 2007  Internet Systems Consortium, Inc. ("ISC")
5 * Copyright (C) 1999-2002  Internet Software Consortium.
6 *
7 * Permission to use, copy, modify, and/or distribute this software for any
8 * purpose with or without fee is hereby granted, provided that the above
9 * copyright notice and this permission notice appear in all copies.
10 *
11 * THE SOFTWARE IS PROVIDED "AS IS" AND ISC DISCLAIMS ALL WARRANTIES WITH
12 * REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
13 * AND FITNESS.  IN NO EVENT SHALL ISC BE LIABLE FOR ANY SPECIAL, DIRECT,
14 * INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
15 * LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
16 * OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
17 * PERFORMANCE OF THIS SOFTWARE.
18 */
19
20/* Id: ratelimiter.c,v 1.25 2007/06/19 23:47:17 tbox Exp  */
21
22/*! \file */
23
24#include <config.h>
25
26#include <isc/mem.h>
27#include <isc/ratelimiter.h>
28#include <isc/task.h>
29#include <isc/time.h>
30#include <isc/timer.h>
31#include <isc/util.h>
32
33typedef enum {
34	isc_ratelimiter_stalled = 0,
35	isc_ratelimiter_ratelimited = 1,
36	isc_ratelimiter_idle = 2,
37	isc_ratelimiter_shuttingdown = 3
38} isc_ratelimiter_state_t;
39
40struct isc_ratelimiter {
41	isc_mem_t *		mctx;
42	isc_mutex_t		lock;
43	int			refs;
44	isc_task_t *		task;
45	isc_timer_t *		timer;
46	isc_interval_t		interval;
47	isc_uint32_t		pertic;
48	isc_ratelimiter_state_t	state;
49	isc_event_t		shutdownevent;
50	ISC_LIST(isc_event_t)	pending;
51};
52
53#define ISC_RATELIMITEREVENT_SHUTDOWN (ISC_EVENTCLASS_RATELIMITER + 1)
54
55static void
56ratelimiter_tick(isc_task_t *task, isc_event_t *event);
57
58static void
59ratelimiter_shutdowncomplete(isc_task_t *task, isc_event_t *event);
60
61isc_result_t
62isc_ratelimiter_create(isc_mem_t *mctx, isc_timermgr_t *timermgr,
63		       isc_task_t *task, isc_ratelimiter_t **ratelimiterp)
64{
65	isc_result_t result;
66	isc_ratelimiter_t *rl;
67	INSIST(ratelimiterp != NULL && *ratelimiterp == NULL);
68
69	rl = isc_mem_get(mctx, sizeof(*rl));
70	if (rl == NULL)
71		return ISC_R_NOMEMORY;
72	rl->mctx = mctx;
73	rl->refs = 1;
74	rl->task = task;
75	isc_interval_set(&rl->interval, 0, 0);
76	rl->timer = NULL;
77	rl->pertic = 1;
78	rl->state = isc_ratelimiter_idle;
79	ISC_LIST_INIT(rl->pending);
80
81	result = isc_mutex_init(&rl->lock);
82	if (result != ISC_R_SUCCESS)
83		goto free_mem;
84	result = isc_timer_create(timermgr, isc_timertype_inactive,
85				  NULL, NULL, rl->task, ratelimiter_tick,
86				  rl, &rl->timer);
87	if (result != ISC_R_SUCCESS)
88		goto free_mutex;
89
90	/*
91	 * Increment the reference count to indicate that we may
92	 * (soon) have events outstanding.
93	 */
94	rl->refs++;
95
96	ISC_EVENT_INIT(&rl->shutdownevent,
97		       sizeof(isc_event_t),
98		       0, NULL, ISC_RATELIMITEREVENT_SHUTDOWN,
99		       ratelimiter_shutdowncomplete, rl, rl, NULL, NULL);
100
101	*ratelimiterp = rl;
102	return (ISC_R_SUCCESS);
103
104free_mutex:
105	DESTROYLOCK(&rl->lock);
106free_mem:
107	isc_mem_put(mctx, rl, sizeof(*rl));
108	return (result);
109}
110
111isc_result_t
112isc_ratelimiter_setinterval(isc_ratelimiter_t *rl, isc_interval_t *interval) {
113	isc_result_t result = ISC_R_SUCCESS;
114	LOCK(&rl->lock);
115	rl->interval = *interval;
116	/*
117	 * If the timer is currently running, change its rate.
118	 */
119        if (rl->state == isc_ratelimiter_ratelimited) {
120		result = isc_timer_reset(rl->timer, isc_timertype_ticker, NULL,
121					 &rl->interval, ISC_FALSE);
122	}
123	UNLOCK(&rl->lock);
124	return (result);
125}
126
127void
128isc_ratelimiter_setpertic(isc_ratelimiter_t *rl, isc_uint32_t pertic) {
129	if (pertic == 0)
130		pertic = 1;
131	rl->pertic = pertic;
132}
133
134isc_result_t
135isc_ratelimiter_enqueue(isc_ratelimiter_t *rl, isc_task_t *task,
136			isc_event_t **eventp)
137{
138	isc_result_t result = ISC_R_SUCCESS;
139	isc_event_t *ev;
140
141	REQUIRE(eventp != NULL && *eventp != NULL);
142	REQUIRE(task != NULL);
143	ev = *eventp;
144	REQUIRE(ev->ev_sender == NULL);
145
146	LOCK(&rl->lock);
147        if (rl->state == isc_ratelimiter_ratelimited ||
148	    rl->state == isc_ratelimiter_stalled) {
149		isc_event_t *ev = *eventp;
150		ev->ev_sender = task;
151                ISC_LIST_APPEND(rl->pending, ev, ev_link);
152		*eventp = NULL;
153        } else if (rl->state == isc_ratelimiter_idle) {
154		result = isc_timer_reset(rl->timer, isc_timertype_ticker, NULL,
155					 &rl->interval, ISC_FALSE);
156		if (result == ISC_R_SUCCESS) {
157			ev->ev_sender = task;
158			rl->state = isc_ratelimiter_ratelimited;
159		}
160	} else {
161		INSIST(rl->state == isc_ratelimiter_shuttingdown);
162		result = ISC_R_SHUTTINGDOWN;
163	}
164	UNLOCK(&rl->lock);
165	if (*eventp != NULL && result == ISC_R_SUCCESS)
166		isc_task_send(task, eventp);
167	return (result);
168}
169
170static void
171ratelimiter_tick(isc_task_t *task, isc_event_t *event) {
172	isc_result_t result = ISC_R_SUCCESS;
173	isc_ratelimiter_t *rl = (isc_ratelimiter_t *)event->ev_arg;
174	isc_event_t *p;
175	isc_uint32_t pertic;
176
177	UNUSED(task);
178
179	isc_event_free(&event);
180
181	pertic = rl->pertic;
182        while (pertic != 0) {
183		pertic--;
184		LOCK(&rl->lock);
185		p = ISC_LIST_HEAD(rl->pending);
186		if (p != NULL) {
187			/*
188			 * There is work to do.  Let's do it after unlocking.
189			 */
190			ISC_LIST_UNLINK(rl->pending, p, ev_link);
191		} else {
192			/*
193			 * No work left to do.  Stop the timer so that we don't
194			 * waste resources by having it fire periodically.
195			 */
196			result = isc_timer_reset(rl->timer,
197						 isc_timertype_inactive,
198						 NULL, NULL, ISC_FALSE);
199			RUNTIME_CHECK(result == ISC_R_SUCCESS);
200			rl->state = isc_ratelimiter_idle;
201			pertic = 0;	/* Force the loop to exit. */
202		}
203		UNLOCK(&rl->lock);
204		if (p != NULL) {
205			isc_task_t *evtask = p->ev_sender;
206			isc_task_send(evtask, &p);
207		}
208		INSIST(p == NULL);
209	}
210}
211
212void
213isc_ratelimiter_shutdown(isc_ratelimiter_t *rl) {
214	isc_event_t *ev;
215	isc_task_t *task;
216	LOCK(&rl->lock);
217	rl->state = isc_ratelimiter_shuttingdown;
218	(void)isc_timer_reset(rl->timer, isc_timertype_inactive,
219			      NULL, NULL, ISC_FALSE);
220	while ((ev = ISC_LIST_HEAD(rl->pending)) != NULL) {
221		ISC_LIST_UNLINK(rl->pending, ev, ev_link);
222		ev->ev_attributes |= ISC_EVENTATTR_CANCELED;
223		task = ev->ev_sender;
224		isc_task_send(task, &ev);
225	}
226	isc_timer_detach(&rl->timer);
227	/*
228	 * Send an event to our task.  The delivery of this event
229	 * indicates that no more timer events will be delivered.
230	 */
231	ev = &rl->shutdownevent;
232	isc_task_send(rl->task, &ev);
233
234	UNLOCK(&rl->lock);
235}
236
237static void
238ratelimiter_shutdowncomplete(isc_task_t *task, isc_event_t *event) {
239	isc_ratelimiter_t *rl = (isc_ratelimiter_t *)event->ev_arg;
240
241	UNUSED(task);
242
243	isc_ratelimiter_detach(&rl);
244}
245
246static void
247ratelimiter_free(isc_ratelimiter_t *rl) {
248	DESTROYLOCK(&rl->lock);
249	isc_mem_put(rl->mctx, rl, sizeof(*rl));
250}
251
252void
253isc_ratelimiter_attach(isc_ratelimiter_t *source, isc_ratelimiter_t **target) {
254	REQUIRE(source != NULL);
255	REQUIRE(target != NULL && *target == NULL);
256
257	LOCK(&source->lock);
258	REQUIRE(source->refs > 0);
259	source->refs++;
260	INSIST(source->refs > 0);
261	UNLOCK(&source->lock);
262	*target = source;
263}
264
265void
266isc_ratelimiter_detach(isc_ratelimiter_t **rlp) {
267	isc_ratelimiter_t *rl = *rlp;
268	isc_boolean_t free_now = ISC_FALSE;
269
270	LOCK(&rl->lock);
271	REQUIRE(rl->refs > 0);
272	rl->refs--;
273	if (rl->refs == 0)
274		free_now = ISC_TRUE;
275	UNLOCK(&rl->lock);
276
277	if (free_now)
278		ratelimiter_free(rl);
279
280	*rlp = NULL;
281}
282
283isc_result_t
284isc_ratelimiter_stall(isc_ratelimiter_t *rl) {
285	isc_result_t result = ISC_R_SUCCESS;
286
287	LOCK(&rl->lock);
288	switch (rl->state) {
289	case isc_ratelimiter_shuttingdown:
290		result = ISC_R_SHUTTINGDOWN;
291		break;
292	case isc_ratelimiter_ratelimited:
293		result = isc_timer_reset(rl->timer, isc_timertype_inactive,
294				 	 NULL, NULL, ISC_FALSE);
295		RUNTIME_CHECK(result == ISC_R_SUCCESS);
296	case isc_ratelimiter_idle:
297	case isc_ratelimiter_stalled:
298		rl->state = isc_ratelimiter_stalled;
299		break;
300	}
301	UNLOCK(&rl->lock);
302	return (result);
303}
304
305isc_result_t
306isc_ratelimiter_release(isc_ratelimiter_t *rl) {
307	isc_result_t result = ISC_R_SUCCESS;
308
309	LOCK(&rl->lock);
310	switch (rl->state) {
311	case isc_ratelimiter_shuttingdown:
312		result = ISC_R_SHUTTINGDOWN;
313		break;
314	case isc_ratelimiter_stalled:
315		if (!ISC_LIST_EMPTY(rl->pending)) {
316			result = isc_timer_reset(rl->timer,
317						 isc_timertype_ticker, NULL,
318						 &rl->interval, ISC_FALSE);
319			if (result == ISC_R_SUCCESS)
320				rl->state = isc_ratelimiter_ratelimited;
321		} else
322			rl->state = isc_ratelimiter_idle;
323		break;
324	case isc_ratelimiter_ratelimited:
325	case isc_ratelimiter_idle:
326		break;
327	}
328	UNLOCK(&rl->lock);
329	return (result);
330}
331