xref: /netbsd-src/external/mpl/bind/dist/lib/isc/quota.c (revision b5c47949a45ac972130c38cf13dfd8afb1f09285)
1 /*	$NetBSD: quota.c,v 1.7 2021/02/19 16:42:19 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 https://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 <stddef.h>
17 
18 #include <isc/atomic.h>
19 #include <isc/quota.h>
20 #include <isc/util.h>
21 
22 #define QUOTA_MAGIC    ISC_MAGIC('Q', 'U', 'O', 'T')
23 #define VALID_QUOTA(p) ISC_MAGIC_VALID(p, QUOTA_MAGIC)
24 
25 #define QUOTA_CB_MAGIC	  ISC_MAGIC('Q', 'T', 'C', 'B')
26 #define VALID_QUOTA_CB(p) ISC_MAGIC_VALID(p, QUOTA_CB_MAGIC)
27 
28 void
29 isc_quota_init(isc_quota_t *quota, unsigned int max) {
30 	atomic_init(&quota->max, max);
31 	atomic_init(&quota->used, 0);
32 	atomic_init(&quota->soft, 0);
33 	atomic_init(&quota->waiting, 0);
34 	ISC_LIST_INIT(quota->cbs);
35 	isc_mutex_init(&quota->cblock);
36 	quota->magic = QUOTA_MAGIC;
37 }
38 
39 void
40 isc_quota_destroy(isc_quota_t *quota) {
41 	REQUIRE(VALID_QUOTA(quota));
42 	quota->magic = 0;
43 
44 	INSIST(atomic_load(&quota->used) == 0);
45 	INSIST(atomic_load(&quota->waiting) == 0);
46 	INSIST(ISC_LIST_EMPTY(quota->cbs));
47 	atomic_store_release(&quota->max, 0);
48 	atomic_store_release(&quota->used, 0);
49 	atomic_store_release(&quota->soft, 0);
50 	isc_mutex_destroy(&quota->cblock);
51 }
52 
53 void
54 isc_quota_soft(isc_quota_t *quota, unsigned int soft) {
55 	REQUIRE(VALID_QUOTA(quota));
56 	atomic_store_release(&quota->soft, soft);
57 }
58 
59 void
60 isc_quota_max(isc_quota_t *quota, unsigned int max) {
61 	REQUIRE(VALID_QUOTA(quota));
62 	atomic_store_release(&quota->max, max);
63 }
64 
65 unsigned int
66 isc_quota_getmax(isc_quota_t *quota) {
67 	REQUIRE(VALID_QUOTA(quota));
68 	return (atomic_load_relaxed(&quota->max));
69 }
70 
71 unsigned int
72 isc_quota_getsoft(isc_quota_t *quota) {
73 	REQUIRE(VALID_QUOTA(quota));
74 	return (atomic_load_relaxed(&quota->soft));
75 }
76 
77 unsigned int
78 isc_quota_getused(isc_quota_t *quota) {
79 	REQUIRE(VALID_QUOTA(quota));
80 	return (atomic_load_relaxed(&quota->used));
81 }
82 
83 static isc_result_t
84 quota_reserve(isc_quota_t *quota) {
85 	isc_result_t result;
86 	uint_fast32_t max = atomic_load_acquire(&quota->max);
87 	uint_fast32_t soft = atomic_load_acquire(&quota->soft);
88 	uint_fast32_t used = atomic_load_acquire(&quota->used);
89 	do {
90 		if (max != 0 && used >= max) {
91 			return (ISC_R_QUOTA);
92 		}
93 		if (soft != 0 && used >= soft) {
94 			result = ISC_R_SOFTQUOTA;
95 		} else {
96 			result = ISC_R_SUCCESS;
97 		}
98 	} while (!atomic_compare_exchange_weak_acq_rel(&quota->used, &used,
99 						       used + 1));
100 	return (result);
101 }
102 
103 /* Must be quota->cbslock locked */
104 static void
105 enqueue(isc_quota_t *quota, isc_quota_cb_t *cb) {
106 	REQUIRE(cb != NULL);
107 	ISC_LIST_ENQUEUE(quota->cbs, cb, link);
108 	atomic_fetch_add_release(&quota->waiting, 1);
109 }
110 
111 /* Must be quota->cbslock locked */
112 static isc_quota_cb_t *
113 dequeue(isc_quota_t *quota) {
114 	isc_quota_cb_t *cb = ISC_LIST_HEAD(quota->cbs);
115 	INSIST(cb != NULL);
116 	ISC_LIST_DEQUEUE(quota->cbs, cb, link);
117 	atomic_fetch_sub_relaxed(&quota->waiting, 1);
118 	return (cb);
119 }
120 
121 static void
122 quota_release(isc_quota_t *quota) {
123 	/*
124 	 * This is opportunistic - we might race with a failing quota_attach_cb
125 	 * and not detect that something is waiting, but eventually someone will
126 	 * be releasing quota and will detect it, so we don't need to worry -
127 	 * and we're saving a lot by not locking cblock every time.
128 	 */
129 
130 	if (atomic_load_acquire(&quota->waiting) > 0) {
131 		isc_quota_cb_t *cb = NULL;
132 		LOCK(&quota->cblock);
133 		if (atomic_load_relaxed(&quota->waiting) > 0) {
134 			cb = dequeue(quota);
135 		}
136 		UNLOCK(&quota->cblock);
137 		if (cb != NULL) {
138 			cb->cb_func(quota, cb->data);
139 			return;
140 		}
141 	}
142 
143 	INSIST(atomic_fetch_sub_release(&quota->used, 1) > 0);
144 }
145 
146 static isc_result_t
147 doattach(isc_quota_t *quota, isc_quota_t **p) {
148 	isc_result_t result;
149 	REQUIRE(p != NULL && *p == NULL);
150 
151 	result = quota_reserve(quota);
152 	if (result == ISC_R_SUCCESS || result == ISC_R_SOFTQUOTA) {
153 		*p = quota;
154 	}
155 
156 	return (result);
157 }
158 
159 isc_result_t
160 isc_quota_attach(isc_quota_t *quota, isc_quota_t **quotap) {
161 	REQUIRE(VALID_QUOTA(quota));
162 	REQUIRE(quotap != NULL && *quotap == NULL);
163 
164 	return (isc_quota_attach_cb(quota, quotap, NULL));
165 }
166 
167 isc_result_t
168 isc_quota_attach_cb(isc_quota_t *quota, isc_quota_t **quotap,
169 		    isc_quota_cb_t *cb) {
170 	REQUIRE(VALID_QUOTA(quota));
171 	REQUIRE(cb == NULL || VALID_QUOTA_CB(cb));
172 	REQUIRE(quotap != NULL && *quotap == NULL);
173 
174 	isc_result_t result = doattach(quota, quotap);
175 	if (result == ISC_R_QUOTA && cb != NULL) {
176 		LOCK(&quota->cblock);
177 		enqueue(quota, cb);
178 		UNLOCK(&quota->cblock);
179 	}
180 	return (result);
181 }
182 
183 void
184 isc_quota_cb_init(isc_quota_cb_t *cb, isc_quota_cb_func_t cb_func, void *data) {
185 	ISC_LINK_INIT(cb, link);
186 	cb->cb_func = cb_func;
187 	cb->data = data;
188 	cb->magic = QUOTA_CB_MAGIC;
189 }
190 
191 void
192 isc_quota_detach(isc_quota_t **quotap) {
193 	REQUIRE(quotap != NULL && VALID_QUOTA(*quotap));
194 	isc_quota_t *quota = *quotap;
195 	*quotap = NULL;
196 
197 	quota_release(quota);
198 }
199