2
0
mirror of https://gitlab.isc.org/isc-projects/bind9 synced 2025-08-30 05:57:52 +00:00

Merge branch 'wpk/quota-callback' into 'master'

isc_quota_attach_cb - quota attach function with a callback.

See merge request isc-projects/bind9!3280
This commit is contained in:
Witold Krecicki 2020-03-30 08:17:31 +00:00
commit 2df610a882
6 changed files with 471 additions and 53 deletions

View File

@ -40,11 +40,23 @@
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
@ -90,41 +102,46 @@ isc_quota_getused(isc_quota_t *quota);
*/
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
*/
isc_result_t
isc_quota_attach_cb(isc_quota_t *quota, isc_quota_t **p, isc_quota_cb_t *cb);
/*%<
*
* 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
*/
void
isc_quota_release(isc_quota_t *quota);
isc_quota_cb_init(isc_quota_cb_t *cb, isc_quota_cb_func_t cb_func, void *data);
/*%<
* Release one unit of quota.
*/
isc_result_t
isc_quota_attach(isc_quota_t *quota, isc_quota_t **p);
/*%<
* Like isc_quota_reserve, and also attaches '*p' to the
* quota if successful (ISC_R_SUCCESS or ISC_R_SOFTQUOTA).
*/
isc_result_t
isc_quota_force(isc_quota_t *quota, isc_quota_t **p);
/*%<
* 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

View File

@ -22,14 +22,20 @@ 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);
}
void
isc_quota_destroy(isc_quota_t *quota) {
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
@ -57,43 +63,77 @@ isc_quota_getused(isc_quota_t *quota) {
return (atomic_load_relaxed(&quota->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(&quota->max);
uint32_t soft = atomic_load_acquire(&quota->soft);
uint32_t used = atomic_fetch_add_relaxed(&quota->used, 1);
if (max == 0 || used < max) {
if (soft == 0 || used < soft) {
result = ISC_R_SUCCESS;
} else {
result = ISC_R_SOFTQUOTA;
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);
}
} else {
INSIST(atomic_fetch_sub_release(&quota->used, 1) > 0);
result = 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);
}
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(&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, 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(&quota->used, 1);
*p = quota;
result = ISC_R_SUCCESS;
}
return (result);
@ -101,17 +141,30 @@ doattach(isc_quota_t *quota, isc_quota_t **p, bool force) {
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(&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;
}
void
isc_quota_detach(isc_quota_t **p) {
INSIST(p != NULL && *p != NULL);
isc_quota_release(*p);
quota_release(*p);
*p = NULL;
}

View File

@ -33,7 +33,7 @@ SRCS = isctest.c aes_test.c buffer_test.c \
counter_test.c crc64_test.c errno_test.c file_test.c hash_test.c \
heap_test.c hmac_test.c ht_test.c lex_test.c \
mem_test.c md_test.c netaddr_test.c parse_test.c pool_test.c \
radix_test.c random_test.c \
quota_test.c radix_test.c random_test.c \
regex_test.c result_test.c safe_test.c siphash_test.c sockaddr_test.c \
socket_test.c socket_test.c symtab_test.c task_test.c \
taskpool_test.c time_test.c timer_test.c
@ -46,7 +46,7 @@ TARGETS = aes_test@EXEEXT@ buffer_test@EXEEXT@ \
ht_test@EXEEXT@ \
lex_test@EXEEXT@ mem_test@EXEEXT@ md_test@EXEEXT@ \
netaddr_test@EXEEXT@ parse_test@EXEEXT@ pool_test@EXEEXT@ \
radix_test@EXEEXT@ \
quota_test@EXEEXT@ radix_test@EXEEXT@ \
random_test@EXEEXT@ regex_test@EXEEXT@ result_test@EXEEXT@ \
safe_test@EXEEXT@ siphash_test@EXEEXT@ sockaddr_test@EXEEXT@ socket_test@EXEEXT@ \
socket_test@EXEEXT@ symtab_test@EXEEXT@ task_test@EXEEXT@ \
@ -134,6 +134,11 @@ pool_test@EXEEXT@: pool_test.@O@ isctest.@O@ ${ISCDEPLIBS}
${LDFLAGS} -o $@ pool_test.@O@ isctest.@O@ \
${ISCLIBS} ${LIBS}
quota_test@EXEEXT@: quota_test.@O@ isctest.@O@ ${ISCDEPLIBS}
${LIBTOOL_MODE_LINK} ${PURIFY} ${CC} ${CFLAGS} \
${LDFLAGS} -o $@ quota_test.@O@ isctest.@O@ \
${ISCLIBS} ${LIBS}
radix_test@EXEEXT@: radix_test.@O@ isctest.@O@ ${ISCDEPLIBS}
${LIBTOOL_MODE_LINK} ${PURIFY} ${CC} ${CFLAGS} \
${LDFLAGS} -o $@ radix_test.@O@ isctest.@O@ \

343
lib/isc/tests/quota_test.c Normal file
View File

@ -0,0 +1,343 @@
/*
* 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(&quota, 100);
assert_int_equal(isc_quota_getmax(&quota), 100);
assert_int_equal(isc_quota_getsoft(&quota), 0);
isc_quota_max(&quota, 50);
isc_quota_soft(&quota, 30);
assert_int_equal(isc_quota_getmax(&quota), 50);
assert_int_equal(isc_quota_getsoft(&quota), 30);
assert_int_equal(isc_quota_getused(&quota), 0);
isc_quota_attach(&quota, &quota2);
assert_int_equal(isc_quota_getused(&quota), 1);
isc_quota_detach(&quota2);
assert_int_equal(isc_quota_getused(&quota), 0);
isc_quota_destroy(&quota);
}
#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(&quota, 100);
for (i = 0; i < 100; i++) {
add_quota(&quota, &quotas[i], ISC_R_SUCCESS, true, i + 1);
}
add_quota(&quota, &quotas[100], ISC_R_QUOTA, false, 100);
assert_int_equal(isc_quota_getused(&quota), 100);
isc_quota_detach(&quotas[0]);
assert_null(quotas[0]);
add_quota(&quota, &quotas[100], ISC_R_SUCCESS, true, 100);
add_quota(&quota, &quotas[101], ISC_R_QUOTA, false, 100);
for (i = 100; i > 0; i--) {
isc_quota_detach(&quotas[i]);
assert_null(quotas[i]);
assert_int_equal(isc_quota_getused(&quota), i - 1);
}
assert_int_equal(isc_quota_getused(&quota), 0);
isc_quota_destroy(&quota);
}
static void
isc_quota_soft_test(void **state) {
isc_quota_t quota;
isc_quota_t *quotas[110];
int i;
UNUSED(state);
isc_quota_init(&quota, 100);
isc_quota_soft(&quota, 50);
for (i = 0; i < 50; i++) {
add_quota(&quota, &quotas[i], ISC_R_SUCCESS, true, i + 1);
}
for (i = 50; i < 100; i++) {
add_quota(&quota, &quotas[i], ISC_R_SOFTQUOTA, true, i + 1);
}
add_quota(&quota, &quotas[i], ISC_R_QUOTA, false, 100);
for (i = 99; i >= 0; i--) {
isc_quota_detach(&quotas[i]);
assert_null(quotas[i]);
assert_int_equal(isc_quota_getused(&quota), i);
}
assert_int_equal(isc_quota_getused(&quota), 0);
isc_quota_destroy(&quota);
}
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(&quota);
}
}
static void
isc_quota_callback_test(void **state) {
isc_result_t result;
isc_quota_t quota;
isc_quota_t *quotas[30];
qp = &quota;
/*
* - 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(&quota, 20);
isc_quota_soft(&quota, 10);
for (i = 0; i < 10; i++) {
quotas[i] = NULL;
isc_quota_cb_init(&cbs[i], callback, &ints[i]);
result = isc_quota_attach_cb(&quota, &quotas[i], &cbs[i]);
assert_int_equal(result, ISC_R_SUCCESS);
assert_ptr_equal(quotas[i], &quota);
assert_int_equal(isc_quota_getused(&quota), 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(&quota, &quotas[i], &cbs[i]);
assert_int_equal(result, ISC_R_SOFTQUOTA);
assert_ptr_equal(quotas[i], &quota);
assert_int_equal(isc_quota_getused(&quota), 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(&quota, &quotas[i], &cbs[i]);
assert_int_equal(result, ISC_R_QUOTA);
assert_ptr_equal(quotas[i], NULL);
assert_int_equal(isc_quota_getused(&quota), 20);
}
assert_int_equal(atomic_load(&cb_calls), 0);
for (i = 0; i < 5; i++) {
isc_quota_detach(&quotas[i]);
assert_null(quotas[i]);
assert_int_equal(isc_quota_getused(&quota), 20);
assert_int_equal(atomic_load(&cb_calls), i + 1);
}
/* That should cause a chain reaction */
isc_quota_detach(&quotas[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(&quotas[i]);
}
for (i = 6; i < 20; i++) {
isc_quota_detach(&quotas[i]);
assert_null(quotas[i]);
assert_int_equal(isc_quota_getused(&quota), 19 - i);
}
assert_int_equal(atomic_load(&cb_calls), 10);
assert_int_equal(isc_quota_getused(&quota), 0);
isc_quota_destroy(&quota);
}
/*
* 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(&quota);
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, &quota,
&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(&quota, 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 = &quota;
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(&quota);
}
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 */

View File

@ -497,16 +497,15 @@ isc_queue_dequeue
isc_queue_destroy
isc_queue_new
isc_quota_attach
isc_quota_attach_cb
isc_quota_cb_init
isc_quota_destroy
isc_quota_detach
isc_quota_force
isc_quota_getmax
isc_quota_getsoft
isc_quota_getused
isc_quota_init
isc_quota_max
isc_quota_release
isc_quota_reserve
isc_quota_soft
isc_radix_create
isc_radix_destroy

View File

@ -2314,6 +2314,7 @@
./lib/isc/tests/netaddr_test.c C 2016,2018,2019,2020
./lib/isc/tests/parse_test.c C 2012,2013,2016,2018,2019,2020
./lib/isc/tests/pool_test.c C 2013,2016,2018,2019,2020
./lib/isc/tests/quota_test.c C 2020
./lib/isc/tests/radix_test.c C 2014,2016,2018,2019,2020
./lib/isc/tests/random_test.c C 2014,2015,2016,2017,2018,2019,2020
./lib/isc/tests/regex_test.c C 2013,2015,2016,2018,2019,2020