xref: /netbsd-src/external/mpl/bind/dist/lib/isc/quota.c (revision bcda20f65a8566e103791ec395f7f499ef322704)
1 /*	$NetBSD: quota.c,v 1.10 2025/01/26 16:25:38 christos Exp $	*/
2 
3 /*
4  * Copyright (C) Internet Systems Consortium, Inc. ("ISC")
5  *
6  * SPDX-License-Identifier: MPL-2.0
7  *
8  * This Source Code Form is subject to the terms of the Mozilla Public
9  * License, v. 2.0. If a copy of the MPL was not distributed with this
10  * file, you can obtain one at https://mozilla.org/MPL/2.0/.
11  *
12  * See the COPYRIGHT file distributed with this work for additional
13  * information regarding copyright ownership.
14  */
15 
16 /*! \file */
17 
18 #include <stddef.h>
19 
20 #include <isc/atomic.h>
21 #include <isc/quota.h>
22 #include <isc/urcu.h>
23 #include <isc/util.h>
24 
25 #define QUOTA_MAGIC    ISC_MAGIC('Q', 'U', 'O', 'T')
26 #define VALID_QUOTA(p) ISC_MAGIC_VALID(p, QUOTA_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 	cds_wfcq_init(&quota->jobs.head, &quota->jobs.tail);
34 	ISC_LINK_INIT(quota, link);
35 	quota->magic = QUOTA_MAGIC;
36 }
37 
38 void
39 isc_quota_soft(isc_quota_t *quota, unsigned int soft) {
40 	REQUIRE(VALID_QUOTA(quota));
41 	atomic_store_relaxed(&quota->soft, soft);
42 }
43 
44 void
45 isc_quota_max(isc_quota_t *quota, unsigned int max) {
46 	REQUIRE(VALID_QUOTA(quota));
47 	atomic_store_relaxed(&quota->max, max);
48 }
49 
50 unsigned int
51 isc_quota_getmax(isc_quota_t *quota) {
52 	REQUIRE(VALID_QUOTA(quota));
53 	return atomic_load_relaxed(&quota->max);
54 }
55 
56 unsigned int
57 isc_quota_getsoft(isc_quota_t *quota) {
58 	REQUIRE(VALID_QUOTA(quota));
59 	return atomic_load_relaxed(&quota->soft);
60 }
61 
62 unsigned int
63 isc_quota_getused(isc_quota_t *quota) {
64 	REQUIRE(VALID_QUOTA(quota));
65 	return atomic_load_relaxed(&quota->used);
66 }
67 
68 void
69 isc_quota_release(isc_quota_t *quota) {
70 	/*
71 	 * We are using the cds_wfcq_dequeue_blocking() variant here that
72 	 * has an internal mutex because we need synchronization on
73 	 * multiple dequeues running from different threads.
74 	 */
75 	struct cds_wfcq_node *node =
76 		cds_wfcq_dequeue_blocking(&quota->jobs.head, &quota->jobs.tail);
77 	if (node == NULL) {
78 		uint_fast32_t used = atomic_fetch_sub_relaxed(&quota->used, 1);
79 		INSIST(used > 0);
80 		return;
81 	}
82 
83 	isc_job_t *job = caa_container_of(node, isc_job_t, wfcq_node);
84 	job->cb(job->cbarg);
85 }
86 
87 isc_result_t
88 isc_quota_acquire_cb(isc_quota_t *quota, isc_job_t *job, isc_job_cb cb,
89 		     void *cbarg) {
90 	REQUIRE(VALID_QUOTA(quota));
91 	REQUIRE(job == NULL || cb != NULL);
92 
93 	uint_fast32_t used = atomic_fetch_add_relaxed(&quota->used, 1);
94 	uint_fast32_t max = atomic_load_relaxed(&quota->max);
95 	if (max != 0 && used >= max) {
96 		(void)atomic_fetch_sub_relaxed(&quota->used, 1);
97 		if (job != NULL) {
98 			job->cb = cb;
99 			job->cbarg = cbarg;
100 			cds_wfcq_node_init(&job->wfcq_node);
101 
102 			/*
103 			 * The cds_wfcq_enqueue() is non-blocking (no internal
104 			 * mutex involved), so it offers a slight advantage.
105 			 */
106 			cds_wfcq_enqueue(&quota->jobs.head, &quota->jobs.tail,
107 					 &job->wfcq_node);
108 		}
109 		return ISC_R_QUOTA;
110 	}
111 
112 	uint_fast32_t soft = atomic_load_relaxed(&quota->soft);
113 	if (soft != 0 && used >= soft) {
114 		return ISC_R_SOFTQUOTA;
115 	}
116 
117 	return ISC_R_SUCCESS;
118 }
119 
120 void
121 isc_quota_destroy(isc_quota_t *quota) {
122 	REQUIRE(VALID_QUOTA(quota));
123 	quota->magic = 0;
124 
125 	INSIST(atomic_load(&quota->used) == 0);
126 	INSIST(cds_wfcq_empty(&quota->jobs.head, &quota->jobs.tail));
127 
128 	cds_wfcq_destroy(&quota->jobs.head, &quota->jobs.tail);
129 }
130