summaryrefslogtreecommitdiffstats
path: root/lib/isc/quota.c
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--lib/isc/quota.c198
1 files changed, 198 insertions, 0 deletions
diff --git a/lib/isc/quota.c b/lib/isc/quota.c
new file mode 100644
index 0000000..02c7176
--- /dev/null
+++ b/lib/isc/quota.c
@@ -0,0 +1,198 @@
+/*
+ * 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(&quota->max, max);
+ atomic_init(&quota->used, 0);
+ atomic_init(&quota->soft, 0);
+ atomic_init(&quota->waiting, 0);
+ ISC_LIST_INIT(quota->cbs);
+ isc_mutex_init(&quota->cblock);
+ quota->magic = QUOTA_MAGIC;
+}
+
+void
+isc_quota_destroy(isc_quota_t *quota) {
+ REQUIRE(VALID_QUOTA(quota));
+ quota->magic = 0;
+
+ INSIST(atomic_load(&quota->used) == 0);
+ INSIST(atomic_load(&quota->waiting) == 0);
+ INSIST(ISC_LIST_EMPTY(quota->cbs));
+ atomic_store_release(&quota->max, 0);
+ atomic_store_release(&quota->used, 0);
+ atomic_store_release(&quota->soft, 0);
+ isc_mutex_destroy(&quota->cblock);
+}
+
+void
+isc_quota_soft(isc_quota_t *quota, unsigned int soft) {
+ REQUIRE(VALID_QUOTA(quota));
+ atomic_store_release(&quota->soft, soft);
+}
+
+void
+isc_quota_max(isc_quota_t *quota, unsigned int max) {
+ REQUIRE(VALID_QUOTA(quota));
+ atomic_store_release(&quota->max, max);
+}
+
+unsigned int
+isc_quota_getmax(isc_quota_t *quota) {
+ REQUIRE(VALID_QUOTA(quota));
+ return (atomic_load_relaxed(&quota->max));
+}
+
+unsigned int
+isc_quota_getsoft(isc_quota_t *quota) {
+ REQUIRE(VALID_QUOTA(quota));
+ return (atomic_load_relaxed(&quota->soft));
+}
+
+unsigned int
+isc_quota_getused(isc_quota_t *quota) {
+ REQUIRE(VALID_QUOTA(quota));
+ return (atomic_load_relaxed(&quota->used));
+}
+
+static isc_result_t
+quota_reserve(isc_quota_t *quota) {
+ isc_result_t result;
+ uint_fast32_t max = atomic_load_acquire(&quota->max);
+ uint_fast32_t soft = atomic_load_acquire(&quota->soft);
+ uint_fast32_t used = atomic_load_acquire(&quota->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(&quota->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(&quota->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(&quota->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(&quota->waiting) > 0) {
+ isc_quota_cb_t *cb = NULL;
+ LOCK(&quota->cblock);
+ if (atomic_load_relaxed(&quota->waiting) > 0) {
+ cb = dequeue(quota);
+ }
+ UNLOCK(&quota->cblock);
+ if (cb != NULL) {
+ cb->cb_func(quota, cb->data);
+ return;
+ }
+ }
+
+ INSIST(atomic_fetch_sub_release(&quota->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(&quota->cblock);
+ enqueue(quota, cb);
+ UNLOCK(&quota->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);
+}