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