diff options
Diffstat (limited to 'lib/isc/quota.c')
| -rw-r--r-- | lib/isc/quota.c | 202 |
1 files changed, 202 insertions, 0 deletions
diff --git a/lib/isc/quota.c b/lib/isc/quota.c new file mode 100644 index 0000000..5465db5 --- /dev/null +++ b/lib/isc/quota.c @@ -0,0 +1,202 @@ +/* + * Copyright (C) Internet Systems Consortium, Inc. ("ISC") + * + * SPDX-License-Identifier: MPL-2.0 + * + * This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, you can obtain one at https://mozilla.org/MPL/2.0/. + * + * See the COPYRIGHT file distributed with this work for additional + * information regarding copyright ownership. + */ + +/*! \file */ + +#include <stddef.h> + +#include <isc/atomic.h> +#include <isc/quota.h> +#include <isc/util.h> + +#define QUOTA_MAGIC ISC_MAGIC('Q', 'U', 'O', 'T') +#define VALID_QUOTA(p) ISC_MAGIC_VALID(p, QUOTA_MAGIC) + +#define QUOTA_CB_MAGIC ISC_MAGIC('Q', 'T', 'C', 'B') +#define VALID_QUOTA_CB(p) ISC_MAGIC_VALID(p, QUOTA_CB_MAGIC) + +void +isc_quota_init(isc_quota_t *quota, unsigned int max) { + atomic_init("a->max, max); + atomic_init("a->used, 0); + atomic_init("a->soft, 0); + atomic_init("a->waiting, 0); + ISC_LIST_INIT(quota->cbs); + isc_mutex_init("a->cblock); + ISC_LINK_INIT(quota, link); + quota->magic = QUOTA_MAGIC; +} + +void +isc_quota_destroy(isc_quota_t *quota) { + REQUIRE(VALID_QUOTA(quota)); + quota->magic = 0; + + INSIST(atomic_load("a->used) == 0); + INSIST(atomic_load("a->waiting) == 0); + INSIST(ISC_LIST_EMPTY(quota->cbs)); + atomic_store_release("a->max, 0); + atomic_store_release("a->used, 0); + atomic_store_release("a->soft, 0); + isc_mutex_destroy("a->cblock); +} + +void +isc_quota_soft(isc_quota_t *quota, unsigned int soft) { + REQUIRE(VALID_QUOTA(quota)); + atomic_store_release("a->soft, soft); +} + +void +isc_quota_max(isc_quota_t *quota, unsigned int max) { + REQUIRE(VALID_QUOTA(quota)); + atomic_store_release("a->max, max); +} + +unsigned int +isc_quota_getmax(isc_quota_t *quota) { + REQUIRE(VALID_QUOTA(quota)); + return (atomic_load_relaxed("a->max)); +} + +unsigned int +isc_quota_getsoft(isc_quota_t *quota) { + REQUIRE(VALID_QUOTA(quota)); + return (atomic_load_relaxed("a->soft)); +} + +unsigned int +isc_quota_getused(isc_quota_t *quota) { + REQUIRE(VALID_QUOTA(quota)); + return (atomic_load_relaxed("a->used)); +} + +static isc_result_t +quota_reserve(isc_quota_t *quota) { + isc_result_t result; + uint_fast32_t max = atomic_load_acquire("a->max); + uint_fast32_t soft = atomic_load_acquire("a->soft); + uint_fast32_t used = atomic_load_acquire("a->used); + do { + if (max != 0 && used >= max) { + return (ISC_R_QUOTA); + } + if (soft != 0 && used >= soft) { + result = ISC_R_SOFTQUOTA; + } else { + result = ISC_R_SUCCESS; + } + } while (!atomic_compare_exchange_weak_acq_rel("a->used, &used, + used + 1)); + return (result); +} + +/* Must be quota->cbslock locked */ +static void +enqueue(isc_quota_t *quota, isc_quota_cb_t *cb) { + REQUIRE(cb != NULL); + ISC_LIST_ENQUEUE(quota->cbs, cb, link); + atomic_fetch_add_release("a->waiting, 1); +} + +/* Must be quota->cbslock locked */ +static isc_quota_cb_t * +dequeue(isc_quota_t *quota) { + isc_quota_cb_t *cb = ISC_LIST_HEAD(quota->cbs); + INSIST(cb != NULL); + ISC_LIST_DEQUEUE(quota->cbs, cb, link); + atomic_fetch_sub_relaxed("a->waiting, 1); + return (cb); +} + +static void +quota_release(isc_quota_t *quota) { + uint_fast32_t used; + + /* + * This is opportunistic - we might race with a failing quota_attach_cb + * and not detect that something is waiting, but eventually someone will + * be releasing quota and will detect it, so we don't need to worry - + * and we're saving a lot by not locking cblock every time. + */ + + if (atomic_load_acquire("a->waiting) > 0) { + isc_quota_cb_t *cb = NULL; + LOCK("a->cblock); + if (atomic_load_relaxed("a->waiting) > 0) { + cb = dequeue(quota); + } + UNLOCK("a->cblock); + if (cb != NULL) { + cb->cb_func(quota, cb->data); + return; + } + } + + used = atomic_fetch_sub_release("a->used, 1); + INSIST(used > 0); +} + +static isc_result_t +doattach(isc_quota_t *quota, isc_quota_t **p) { + isc_result_t result; + REQUIRE(p != NULL && *p == NULL); + + result = quota_reserve(quota); + if (result == ISC_R_SUCCESS || result == ISC_R_SOFTQUOTA) { + *p = quota; + } + + return (result); +} + +isc_result_t +isc_quota_attach(isc_quota_t *quota, isc_quota_t **quotap) { + REQUIRE(VALID_QUOTA(quota)); + REQUIRE(quotap != NULL && *quotap == NULL); + + return (isc_quota_attach_cb(quota, quotap, NULL)); +} + +isc_result_t +isc_quota_attach_cb(isc_quota_t *quota, isc_quota_t **quotap, + isc_quota_cb_t *cb) { + REQUIRE(VALID_QUOTA(quota)); + REQUIRE(cb == NULL || VALID_QUOTA_CB(cb)); + REQUIRE(quotap != NULL && *quotap == NULL); + + isc_result_t result = doattach(quota, quotap); + if (result == ISC_R_QUOTA && cb != NULL) { + LOCK("a->cblock); + enqueue(quota, cb); + UNLOCK("a->cblock); + } + return (result); +} + +void +isc_quota_cb_init(isc_quota_cb_t *cb, isc_quota_cb_func_t cb_func, void *data) { + ISC_LINK_INIT(cb, link); + cb->cb_func = cb_func; + cb->data = data; + cb->magic = QUOTA_CB_MAGIC; +} + +void +isc_quota_detach(isc_quota_t **quotap) { + REQUIRE(quotap != NULL && VALID_QUOTA(*quotap)); + isc_quota_t *quota = *quotap; + *quotap = NULL; + + quota_release(quota); +} |