Commit d151a10f authored by Witold Krecicki's avatar Witold Krecicki Committed by Witold Krecicki
Browse files

Add a quota attach function with a callback, some code cleanups.

We introduce a isc_quota_attach_cb function - if ISC_R_QUOTA is returned
at the time the function is called, then a callback will be called when
there's quota available (with quota already attached). The callbacks are
organized as a LIFO queue in the quota structure.
It's needed for TCP client quota -  with old networking code we had one
single place where tcp clients quota was processed so we could resume
accepting when the we had spare slots, but it's gone with netmgr - now
we need to notify the listener/accepter that there's quota available so
that it can resume accepting.

Remove unused isc_quota_force() function.

The isc_quote_reserve and isc_quota_release were used only internally
from the quota.c and the tests.  We should not expose API we are not
using.
parent 5826c85a
Pipeline #37909 passed with stages
in 29 minutes and 15 seconds
......@@ -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
*/
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
......
......@@ -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 {
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;
}
} else {
INSIST(atomic_fetch_sub_release(&quota->used, 1) > 0);
result = ISC_R_QUOTA;
}
} 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;
}
......@@ -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@ \
......
/*
* 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");