quota.c 4.13 KB
Newer Older
1
/*
2
 * Copyright (C) Internet Systems Consortium, Inc. ("ISC")
3
 *
4 5 6
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
7 8 9
 *
 * See the COPYRIGHT file distributed with this work for additional
 * information regarding copyright ownership.
10
 */
11

12
/*! \file */
David Lawrence's avatar
David Lawrence committed
13

Brian Wellington's avatar
Brian Wellington committed
14 15
#include <stddef.h>

16
#include <isc/atomic.h>
17 18 19
#include <isc/quota.h>
#include <isc/util.h>

20
void
Evan Hunt's avatar
Evan Hunt committed
21
isc_quota_init(isc_quota_t *quota, unsigned int max) {
22 23 24
	atomic_init(&quota->max, max);
	atomic_init(&quota->used, 0);
	atomic_init(&quota->soft, 0);
25 26 27
	atomic_init(&quota->waiting, 0);
	ISC_LIST_INIT(quota->cbs);
	isc_mutex_init(&quota->cblock);
28 29 30
}

void
Evan Hunt's avatar
Evan Hunt committed
31
isc_quota_destroy(isc_quota_t *quota) {
32
	INSIST(atomic_load(&quota->used) == 0);
33 34
	INSIST(atomic_load(&quota->waiting) == 0);
	INSIST(ISC_LIST_EMPTY(quota->cbs));
35 36 37
	atomic_store_release(&quota->max, 0);
	atomic_store_release(&quota->used, 0);
	atomic_store_release(&quota->soft, 0);
38
	isc_mutex_destroy(&quota->cblock);
39 40
}

41
void
Evan Hunt's avatar
Evan Hunt committed
42
isc_quota_soft(isc_quota_t *quota, unsigned int soft) {
43
	atomic_store_release(&quota->soft, soft);
44 45 46
}

void
Evan Hunt's avatar
Evan Hunt committed
47
isc_quota_max(isc_quota_t *quota, unsigned int max) {
48
	atomic_store_release(&quota->max, max);
49 50 51
}

unsigned int
Evan Hunt's avatar
Evan Hunt committed
52
isc_quota_getmax(isc_quota_t *quota) {
53
	return (atomic_load_relaxed(&quota->max));
54 55 56
}

unsigned int
Evan Hunt's avatar
Evan Hunt committed
57
isc_quota_getsoft(isc_quota_t *quota) {
58
	return (atomic_load_relaxed(&quota->soft));
59 60 61
}

unsigned int
Evan Hunt's avatar
Evan Hunt committed
62
isc_quota_getused(isc_quota_t *quota) {
63
	return (atomic_load_relaxed(&quota->used));
64 65
}

66 67
static isc_result_t
quota_reserve(isc_quota_t *quota) {
68
	isc_result_t result;
69 70 71 72 73 74 75 76
	uint_fast32_t max = atomic_load_acquire(&quota->max);
	uint_fast32_t soft = atomic_load_acquire(&quota->soft);
	uint_fast32_t used = atomic_load_acquire(&quota->used);
	do {
		if (max != 0 && used >= max) {
			return (ISC_R_QUOTA);
		}
		if (soft != 0 && used >= soft) {
77
			result = ISC_R_SOFTQUOTA;
78 79
		} else {
			result = ISC_R_SUCCESS;
80
		}
81 82
	} while (!atomic_compare_exchange_weak_acq_rel(&quota->used, &used,
						       used + 1));
83 84 85
	return (result);
}

86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121
/* Must be quota->cbslock locked */
static void
enqueue(isc_quota_t *quota, isc_quota_cb_t *cb) {
	REQUIRE(cb != NULL);
	ISC_LIST_ENQUEUE(quota->cbs, cb, link);
	atomic_fetch_add_release(&quota->waiting, 1);
}

/* Must be quota->cbslock locked */
static isc_quota_cb_t *
dequeue(isc_quota_t *quota) {
	isc_quota_cb_t *cb = ISC_LIST_HEAD(quota->cbs);
	INSIST(cb != NULL);
	ISC_LIST_DEQUEUE(quota->cbs, cb, link);
	atomic_fetch_sub_relaxed(&quota->waiting, 1);
	return (cb);
}

static void
quota_release(isc_quota_t *quota) {
	/*
	 * This is opportunistic - we might race with a failing quota_attach_cb
	 * and not detect that something is waiting, but eventually someone will
	 * be releasing quota and will detect it, so we don't need to worry -
	 * and we're saving a lot by not locking cblock every time.
	 */

	if (atomic_load_acquire(&quota->waiting) > 0) {
		isc_quota_cb_t *cb = NULL;
		LOCK(&quota->cblock);
		if (atomic_load_relaxed(&quota->waiting) > 0) {
			cb = dequeue(quota);
		}
		UNLOCK(&quota->cblock);
		if (cb != NULL) {
			cb->cb_func(quota, cb->data);
122 123
			cb->cb_func = NULL;
			cb->data = NULL;
124 125 126 127
			return;
		}
	}

128
	INSIST(atomic_fetch_sub_release(&quota->used, 1) > 0);
129 130
}

131
static isc_result_t
132
doattach(isc_quota_t *quota, isc_quota_t **p) {
133
	isc_result_t result;
134 135
	REQUIRE(p != NULL && *p == NULL);

136
	result = quota_reserve(quota);
137 138 139 140
	if (result == ISC_R_SUCCESS || result == ISC_R_SOFTQUOTA) {
		*p = quota;
	}

141
	return (result);
142
}
143

144
isc_result_t
Evan Hunt's avatar
Evan Hunt committed
145
isc_quota_attach(isc_quota_t *quota, isc_quota_t **p) {
146
	return (isc_quota_attach_cb(quota, p, NULL));
147 148 149
}

isc_result_t
150 151 152 153 154 155 156 157 158 159 160 161 162 163 164
isc_quota_attach_cb(isc_quota_t *quota, isc_quota_t **p, isc_quota_cb_t *cb) {
	isc_result_t result = doattach(quota, p);
	if (result == ISC_R_QUOTA && cb != NULL) {
		LOCK(&quota->cblock);
		enqueue(quota, cb);
		UNLOCK(&quota->cblock);
	}
	return (result);
}

void
isc_quota_cb_init(isc_quota_cb_t *cb, isc_quota_cb_func_t cb_func, void *data) {
	ISC_LINK_INIT(cb, link);
	cb->cb_func = cb_func;
	cb->data = data;
165 166
}

167
void
Evan Hunt's avatar
Evan Hunt committed
168
isc_quota_detach(isc_quota_t **p) {
169
	INSIST(p != NULL && *p != NULL);
170
	quota_release(*p);
171 172
	*p = NULL;
}