ISC_LANG_BEGINDECLS
+/*% isc_quota_cb - quota callback structure */
+typedef struct isc_quota_cb isc_quota_cb_t;
+typedef void (*isc_quota_cb_func_t)(isc_quota_t *quota, void *data);
+struct isc_quota_cb {
+ isc_quota_cb_func_t cb_func;
+ void * data;
+ ISC_LINK(isc_quota_cb_t) link;
+};
+
/*% isc_quota structure */
struct isc_quota {
atomic_uint_fast32_t max;
atomic_uint_fast32_t used;
atomic_uint_fast32_t soft;
+ atomic_uint_fast32_t waiting;
+ isc_mutex_t cblock;
+ ISC_LIST(isc_quota_cb_t) cbs;
};
void
*/
isc_result_t
-isc_quota_reserve(isc_quota_t *quota);
+isc_quota_attach(isc_quota_t *quota, isc_quota_t **p);
/*%<
- * Attempt to reserve one unit of 'quota'.
+ *
+ * Attempt to reserve one unit of 'quota', and also attaches '*p' to the quota
+ * if successful (ISC_R_SUCCESS or ISC_R_SOFTQUOTA).
*
* Returns:
- * \li #ISC_R_SUCCESS Success
+ * \li #ISC_R_SUCCESS Success
* \li #ISC_R_SOFTQUOTA Success soft quota reached
* \li #ISC_R_QUOTA Quota is full
*/
-void
-isc_quota_release(isc_quota_t *quota);
-/*%<
- * Release one unit of quota.
- */
-
isc_result_t
-isc_quota_attach(isc_quota_t *quota, isc_quota_t **p);
+isc_quota_attach_cb(isc_quota_t *quota, isc_quota_t **p, isc_quota_cb_t *cb);
/*%<
- * Like isc_quota_reserve, and also attaches '*p' to the
- * quota if successful (ISC_R_SUCCESS or ISC_R_SOFTQUOTA).
+ *
+ * Like isc_quota_attach(), but if there's no quota left then cb->cb_func will
+ * be called when we are attached to quota.
+ * Note: It's the callee responsibility to make sure that we don't end up with
+ * extremely huge number of callbacks waiting - making it easy to create a
+ * resource exhaustion attack. For example in case of TCP listening we simply
+ * don't accept new connections - so the number of callbacks waiting in the
+ * queue is limited by listen() backlog.
+ *
+ * Returns:
+ * \li #ISC_R_SUCCESS Success
+ * \li #ISC_R_SOFTQUOTA Success soft quota reached
+ * \li #ISC_R_QUOTA Quota is full
*/
-isc_result_t
-isc_quota_force(isc_quota_t *quota, isc_quota_t **p);
+void
+isc_quota_cb_init(isc_quota_cb_t *cb, isc_quota_cb_func_t cb_func, void *data);
/*%<
- * Like isc_quota_attach, but will attach '*p' to the quota
- * even if the hard quota has been exceeded.
+ * Initialize isc_quota_cb_t - setup the list, set the callback and data.
*/
void
isc_quota_detach(isc_quota_t **p);
/*%<
- * Like isc_quota_release, and also detaches '*p' from the
- * quota.
+ * Release one unit of quota, and also detaches '*p' from the quota.
*/
ISC_LANG_ENDDECLS
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);
}
void
isc_quota_destroy(isc_quota_t *quota) {
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
return (atomic_load_relaxed("a->used));
}
-isc_result_t
-isc_quota_reserve(isc_quota_t *quota) {
+static isc_result_t
+quota_reserve(isc_quota_t *quota) {
isc_result_t result;
- uint32_t max = atomic_load_acquire("a->max);
- uint32_t soft = atomic_load_acquire("a->soft);
- uint32_t used = atomic_fetch_add_relaxed("a->used, 1);
- if (max == 0 || used < max) {
- if (soft == 0 || used < soft) {
- result = ISC_R_SUCCESS;
- } else {
+ 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;
}
- } else {
- INSIST(atomic_fetch_sub_release("a->used, 1) > 0);
- result = ISC_R_QUOTA;
- }
+ } while (!atomic_compare_exchange_weak_acq_rel("a->used, &used,
+ used + 1));
return (result);
}
-void
-isc_quota_release(isc_quota_t *quota) {
+/* 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, bool force) {
+doattach(isc_quota_t *quota, isc_quota_t **p) {
isc_result_t result;
REQUIRE(p != NULL && *p == NULL);
- result = isc_quota_reserve(quota);
+ result = quota_reserve(quota);
if (result == ISC_R_SUCCESS || result == ISC_R_SOFTQUOTA) {
*p = quota;
- } else if (result == ISC_R_QUOTA && force) {
- /* attach anyway */
- atomic_fetch_add_relaxed("a->used, 1);
- *p = quota;
- result = ISC_R_SUCCESS;
}
return (result);
isc_result_t
isc_quota_attach(isc_quota_t *quota, isc_quota_t **p) {
- return (doattach(quota, p, false));
+ return (isc_quota_attach_cb(quota, p, NULL));
}
isc_result_t
-isc_quota_force(isc_quota_t *quota, isc_quota_t **p) {
- return (doattach(quota, p, true));
+isc_quota_attach_cb(isc_quota_t *quota, isc_quota_t **p, isc_quota_cb_t *cb) {
+ isc_result_t result = doattach(quota, p);
+ 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;
}
void
isc_quota_detach(isc_quota_t **p) {
INSIST(p != NULL && *p != NULL);
- isc_quota_release(*p);
+ quota_release(*p);
*p = NULL;
}
--- /dev/null
+/*
+ * Copyright (C) Internet Systems Consortium, Inc. ("ISC")
+ *
+ * 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 http://mozilla.org/MPL/2.0/.
+ *
+ * See the COPYRIGHT file distributed with this work for additional
+ * information regarding copyright ownership.
+ */
+
+#if HAVE_CMOCKA
+
+#include <sched.h> /* IWYU pragma: keep */
+#include <setjmp.h>
+#include <stdarg.h>
+#include <stddef.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+
+#define UNIT_TESTING
+#include <cmocka.h>
+
+#include <isc/quota.h>
+#include <isc/result.h>
+#include <isc/thread.h>
+#include <isc/util.h>
+
+static void
+isc_quota_get_set_test(void **state) {
+ UNUSED(state);
+ isc_quota_t quota;
+ isc_quota_t *quota2 = NULL;
+ isc_quota_init("a, 100);
+
+ assert_int_equal(isc_quota_getmax("a), 100);
+ assert_int_equal(isc_quota_getsoft("a), 0);
+
+ isc_quota_max("a, 50);
+ isc_quota_soft("a, 30);
+
+ assert_int_equal(isc_quota_getmax("a), 50);
+ assert_int_equal(isc_quota_getsoft("a), 30);
+
+ assert_int_equal(isc_quota_getused("a), 0);
+ isc_quota_attach("a, "a2);
+ assert_int_equal(isc_quota_getused("a), 1);
+ isc_quota_detach("a2);
+ assert_int_equal(isc_quota_getused("a), 0);
+ isc_quota_destroy("a);
+}
+
+#define add_quota(quota, quotasp, exp, attached, exp_used) \
+ { \
+ *quotasp = NULL; \
+ isc_result_t result = isc_quota_attach(quota, quotasp); \
+ assert_int_equal(result, exp); \
+ if (attached) { \
+ assert_ptr_equal(*quotasp, quota); \
+ } else { \
+ assert_null(*quotasp); \
+ } \
+ assert_int_equal(isc_quota_getused(quota), exp_used); \
+ }
+
+static void
+isc_quota_hard_test(void **state) {
+ isc_quota_t quota;
+ isc_quota_t *quotas[110];
+ int i;
+ UNUSED(state);
+
+ isc_quota_init("a, 100);
+
+ for (i = 0; i < 100; i++) {
+ add_quota("a, "as[i], ISC_R_SUCCESS, true, i + 1);
+ }
+
+ add_quota("a, "as[100], ISC_R_QUOTA, false, 100);
+
+ assert_int_equal(isc_quota_getused("a), 100);
+
+ isc_quota_detach("as[0]);
+ assert_null(quotas[0]);
+
+ add_quota("a, "as[100], ISC_R_SUCCESS, true, 100);
+ add_quota("a, "as[101], ISC_R_QUOTA, false, 100);
+
+ for (i = 100; i > 0; i--) {
+ isc_quota_detach("as[i]);
+ assert_null(quotas[i]);
+ assert_int_equal(isc_quota_getused("a), i - 1);
+ }
+ assert_int_equal(isc_quota_getused("a), 0);
+ isc_quota_destroy("a);
+}
+
+static void
+isc_quota_soft_test(void **state) {
+ isc_quota_t quota;
+ isc_quota_t *quotas[110];
+ int i;
+ UNUSED(state);
+
+ isc_quota_init("a, 100);
+ isc_quota_soft("a, 50);
+
+ for (i = 0; i < 50; i++) {
+ add_quota("a, "as[i], ISC_R_SUCCESS, true, i + 1);
+ }
+ for (i = 50; i < 100; i++) {
+ add_quota("a, "as[i], ISC_R_SOFTQUOTA, true, i + 1);
+ }
+
+ add_quota("a, "as[i], ISC_R_QUOTA, false, 100);
+
+ for (i = 99; i >= 0; i--) {
+ isc_quota_detach("as[i]);
+ assert_null(quotas[i]);
+ assert_int_equal(isc_quota_getused("a), i);
+ }
+ assert_int_equal(isc_quota_getused("a), 0);
+ isc_quota_destroy("a);
+}
+
+static atomic_uint_fast32_t cb_calls = ATOMIC_VAR_INIT(0);
+static isc_quota_cb_t cbs[30];
+static isc_quota_t *qp;
+
+static void
+callback(isc_quota_t *quota, void *data) {
+ int val = *(int *)data;
+ /* Callback is not called if we get the quota directly */
+ assert_int_not_equal(val, -1);
+
+ /* We get the proper quota pointer */
+ assert_ptr_equal(quota, qp);
+
+ /* Verify that the callbacks are called in order */
+ int v = atomic_fetch_add_relaxed(&cb_calls, 1);
+ assert_int_equal(v, val);
+
+ /*
+ * First 5 will be detached by the test function,
+ * for the last 5 - do a 'chain detach'.
+ */
+ if (v >= 5) {
+ isc_quota_detach("a);
+ }
+}
+
+static void
+isc_quota_callback_test(void **state) {
+ isc_result_t result;
+ isc_quota_t quota;
+ isc_quota_t *quotas[30];
+ qp = "a;
+ /*
+ * - 10 calls that end with SUCCESS
+ * - 10 calls that end with SOFTQUOTA
+ * - 10 callbacks
+ */
+ int ints[] = { -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
+ -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
+ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 };
+ int i;
+ UNUSED(state);
+
+ isc_quota_init("a, 20);
+ isc_quota_soft("a, 10);
+
+ for (i = 0; i < 10; i++) {
+ quotas[i] = NULL;
+ isc_quota_cb_init(&cbs[i], callback, &ints[i]);
+ result = isc_quota_attach_cb("a, "as[i], &cbs[i]);
+ assert_int_equal(result, ISC_R_SUCCESS);
+ assert_ptr_equal(quotas[i], "a);
+ assert_int_equal(isc_quota_getused("a), i + 1);
+ }
+ for (i = 10; i < 20; i++) {
+ quotas[i] = NULL;
+ isc_quota_cb_init(&cbs[i], callback, &ints[i]);
+ result = isc_quota_attach_cb("a, "as[i], &cbs[i]);
+ assert_int_equal(result, ISC_R_SOFTQUOTA);
+ assert_ptr_equal(quotas[i], "a);
+ assert_int_equal(isc_quota_getused("a), i + 1);
+ }
+
+ for (i = 20; i < 30; i++) {
+ quotas[i] = NULL;
+ isc_quota_cb_init(&cbs[i], callback, &ints[i]);
+ result = isc_quota_attach_cb("a, "as[i], &cbs[i]);
+ assert_int_equal(result, ISC_R_QUOTA);
+ assert_ptr_equal(quotas[i], NULL);
+ assert_int_equal(isc_quota_getused("a), 20);
+ }
+ assert_int_equal(atomic_load(&cb_calls), 0);
+
+ for (i = 0; i < 5; i++) {
+ isc_quota_detach("as[i]);
+ assert_null(quotas[i]);
+ assert_int_equal(isc_quota_getused("a), 20);
+ assert_int_equal(atomic_load(&cb_calls), i + 1);
+ }
+ /* That should cause a chain reaction */
+ isc_quota_detach("as[5]);
+ assert_int_equal(atomic_load(&cb_calls), 10);
+
+ /* Release the quotas that we did not released in the callback */
+ for (i = 0; i < 5; i++) {
+ isc_quota_detach("as[i]);
+ }
+
+ for (i = 6; i < 20; i++) {
+ isc_quota_detach("as[i]);
+ assert_null(quotas[i]);
+ assert_int_equal(isc_quota_getused("a), 19 - i);
+ }
+ assert_int_equal(atomic_load(&cb_calls), 10);
+
+ assert_int_equal(isc_quota_getused("a), 0);
+ isc_quota_destroy("a);
+}
+
+/*
+ * Multithreaded quota callback test:
+ * - quota set to 100
+ * - 10 threads, each trying to get 100 quotas.
+ * - creates a separate thread to release it after 10ms
+ */
+
+typedef struct qthreadinfo {
+ atomic_uint_fast32_t direct;
+ atomic_uint_fast32_t callback;
+ isc_quota_t *quota;
+ isc_quota_cb_t callbacks[100];
+} qthreadinfo_t;
+
+static atomic_uint_fast32_t g_tnum = ATOMIC_VAR_INIT(0);
+/* at most 10 * 100 quota_detach threads */
+isc_thread_t g_threads[10 * 100];
+
+static void *
+quota_detach(void *quotap) {
+ isc_quota_t *quota = (isc_quota_t *)quotap;
+ usleep(10000);
+ isc_quota_detach("a);
+ return ((isc_threadresult_t)0);
+}
+
+static void
+quota_callback(isc_quota_t *quota, void *data) {
+ qthreadinfo_t *qti = (qthreadinfo_t *)data;
+ atomic_fetch_add_relaxed(&qti->callback, 1);
+ int tnum = atomic_fetch_add_relaxed(&g_tnum, 1);
+ isc_thread_create(quota_detach, quota, &g_threads[tnum]);
+}
+
+static isc_threadresult_t
+quota_thread(void *qtip) {
+ qthreadinfo_t *qti = (qthreadinfo_t *)qtip;
+ for (int i = 0; i < 100; i++) {
+ isc_quota_cb_init(&qti->callbacks[i], quota_callback, qti);
+ isc_quota_t *quota = NULL;
+ isc_result_t result = isc_quota_attach_cb(qti->quota, "a,
+ &qti->callbacks[i]);
+ if (result == ISC_R_SUCCESS) {
+ atomic_fetch_add_relaxed(&qti->direct, 1);
+ int tnum = atomic_fetch_add_relaxed(&g_tnum, 1);
+ isc_thread_create(quota_detach, quota,
+ &g_threads[tnum]);
+ }
+ }
+ return ((isc_threadresult_t)0);
+}
+
+static void
+isc_quota_callback_mt_test(void **state) {
+ UNUSED(state);
+ isc_quota_t quota;
+ int i;
+
+ isc_quota_init("a, 100);
+ static qthreadinfo_t qtis[10];
+ isc_thread_t threads[10];
+ for (i = 0; i < 10; i++) {
+ atomic_init(&qtis[i].direct, 0);
+ atomic_init(&qtis[i].callback, 0);
+ qtis[i].quota = "a;
+ isc_thread_create(quota_thread, &qtis[i], &threads[i]);
+ }
+ for (i = 0; i < 10; i++) {
+ isc_thread_join(threads[i], NULL);
+ }
+
+ for (i = 0; i < (int)atomic_load(&g_tnum); i++) {
+ isc_thread_join(g_threads[i], NULL);
+ }
+ int direct = 0, callback = 0;
+
+ for (i = 0; i < 10; i++) {
+ direct += atomic_load(&qtis[i].direct);
+ callback += atomic_load(&qtis[i].callback);
+ }
+ /* Total quota gained must be 10 threads * 100 tries */
+ assert_int_equal(direct + callback, 10 * 100);
+ /*
+ * At least 100 must be direct, the rest is virtually random:
+ * - in a regular run I'm constantly getting 100:900 ratio
+ * - under rr - usually around ~120:880
+ * - under rr -h - 1000:0
+ */
+ assert_true(direct >= 100);
+
+ isc_quota_destroy("a);
+}
+
+int
+main(void) {
+ const struct CMUnitTest tests[] = {
+ cmocka_unit_test(isc_quota_get_set_test),
+ cmocka_unit_test(isc_quota_hard_test),
+ cmocka_unit_test(isc_quota_soft_test),
+ cmocka_unit_test(isc_quota_callback_test),
+ cmocka_unit_test(isc_quota_callback_mt_test),
+ };
+
+ return (cmocka_run_group_tests(tests, NULL, NULL));
+}
+
+#else /* HAVE_CMOCKA */
+
+#include <stdio.h>
+
+int
+main(void) {
+ printf("1..0 # Skipped: cmocka not available\n");
+ return (0);
+}
+
+#endif /* if HAVE_CMOCKA */
projects / thirdparty / bind9.git / commitdiff
