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