/* $NetBSD: quota.c,v 1.1.2.2 2024/02/24 13:07:21 martin Exp $ */ /* * 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 #include #include #include #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); 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) { /* * 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; } } INSIST(atomic_fetch_sub_release("a->used, 1) > 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); }