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("a->max, max); 31 atomic_init("a->used, 0); 32 atomic_init("a->soft, 0); 33 atomic_init("a->waiting, 0); 34 ISC_LIST_INIT(quota->cbs); 35 isc_mutex_init("a->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("a->used) == 0); 45 INSIST(atomic_load("a->waiting) == 0); 46 INSIST(ISC_LIST_EMPTY(quota->cbs)); 47 atomic_store_release("a->max, 0); 48 atomic_store_release("a->used, 0); 49 atomic_store_release("a->soft, 0); 50 isc_mutex_destroy("a->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("a->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("a->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("a->max)); 69 } 70 71 unsigned int 72 isc_quota_getsoft(isc_quota_t *quota) { 73 REQUIRE(VALID_QUOTA(quota)); 74 return (atomic_load_relaxed("a->soft)); 75 } 76 77 unsigned int 78 isc_quota_getused(isc_quota_t *quota) { 79 REQUIRE(VALID_QUOTA(quota)); 80 return (atomic_load_relaxed("a->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("a->max); 87 uint_fast32_t soft = atomic_load_acquire("a->soft); 88 uint_fast32_t used = atomic_load_acquire("a->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("a->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("a->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("a->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("a->waiting) > 0) { 131 isc_quota_cb_t *cb = NULL; 132 LOCK("a->cblock); 133 if (atomic_load_relaxed("a->waiting) > 0) { 134 cb = dequeue(quota); 135 } 136 UNLOCK("a->cblock); 137 if (cb != NULL) { 138 cb->cb_func(quota, cb->data); 139 return; 140 } 141 } 142 143 INSIST(atomic_fetch_sub_release("a->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("a->cblock); 177 enqueue(quota, cb); 178 UNLOCK("a->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