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