commit 234923a: [Feature] Allow exception when choosing upstream
Vsevolod Stakhov
vsevolod at highsecure.ru
Fri Nov 8 17:49:08 UTC 2019
Author: Vsevolod Stakhov
Date: 2019-11-08 17:46:26 +0000
URL: https://github.com/rspamd/rspamd/commit/234923a5ff0e7ddef989cc4172d3790cc8b7dd12
[Feature] Allow exception when choosing upstream
---
src/libutil/upstream.c | 75 +++++++++++++++++++++++++++++++++++---------------
src/libutil/upstream.h | 11 ++++++++
2 files changed, 64 insertions(+), 22 deletions(-)
diff --git a/src/libutil/upstream.c b/src/libutil/upstream.c
index 3589c4baf..0e0c3707b 100644
--- a/src/libutil/upstream.c
+++ b/src/libutil/upstream.c
@@ -1310,18 +1310,30 @@ rspamd_upstream_restore_cb (gpointer elt, gpointer ls)
}
static struct upstream*
-rspamd_upstream_get_random (struct upstream_list *ups)
+rspamd_upstream_get_random (struct upstream_list *ups,
+ struct upstream *except)
{
- guint idx = ottery_rand_range (ups->alive->len - 1);
+ for (;;) {
+ guint idx = ottery_rand_range (ups->alive->len - 1);
+ struct upstream *up;
- return g_ptr_array_index (ups->alive, idx);
+ up = g_ptr_array_index (ups->alive, idx);
+
+ if (except && up == except) {
+ continue;
+ }
+
+ return up;
+ }
}
static struct upstream*
-rspamd_upstream_get_round_robin (struct upstream_list *ups, gboolean use_cur)
+rspamd_upstream_get_round_robin (struct upstream_list *ups,
+ struct upstream *except,
+ gboolean use_cur)
{
guint max_weight = 0, min_checked = G_MAXUINT;
- struct upstream *up, *selected = NULL, *min_checked_sel = NULL;
+ struct upstream *up = NULL, *selected = NULL, *min_checked_sel = NULL;
guint i;
/* Select upstream with the maximum cur_weight */
@@ -1329,6 +1341,11 @@ rspamd_upstream_get_round_robin (struct upstream_list *ups, gboolean use_cur)
for (i = 0; i < ups->alive->len; i ++) {
up = g_ptr_array_index (ups->alive, i);
+
+ if (except != NULL && up == except) {
+ continue;
+ }
+
if (use_cur) {
if (up->cur_weight > max_weight) {
selected = up;
@@ -1396,18 +1413,15 @@ rspamd_consistent_hash (guint64 key, guint32 nbuckets)
}
static struct upstream*
-rspamd_upstream_get_hashed (struct upstream_list *ups, const guint8 *key, guint keylen)
+rspamd_upstream_get_hashed (struct upstream_list *ups,
+ struct upstream *except,
+ const guint8 *key, guint keylen)
{
guint64 k;
guint32 idx;
static const guint max_tries = 20;
struct upstream *up = NULL;
- if (ups->alive->len == 1) {
- /* Fast path */
- return g_ptr_array_index (ups->alive, 0);
- }
-
/* Generate 64 bits input key */
k = rspamd_cryptobox_fast_hash_specific (RSPAMD_CRYPTOBOX_XXHASH64,
key, keylen, ups->hash_seed);
@@ -1420,8 +1434,8 @@ rspamd_upstream_get_hashed (struct upstream_list *ups, const guint8 *key, guint
idx = rspamd_consistent_hash (k, ups->ups->len);
up = g_ptr_array_index (ups->ups, idx);
- if (up->active_idx < 0) {
- /* Found inactive upstream */
+ if (up->active_idx < 0 || (except != NULL && up == except)) {
+ /* Found inactive or excluded upstream */
k = mum_hash_step (k, ups->hash_seed);
}
else {
@@ -1435,7 +1449,7 @@ rspamd_upstream_get_hashed (struct upstream_list *ups, const guint8 *key, guint
}
/* We failed to find any active upstream */
- up = rspamd_upstream_get_random (ups);
+ up = rspamd_upstream_get_random (ups, except);
msg_info ("failed to find hashed upstream for %s, fallback to random: %s",
ups->ups_line, up->name);
@@ -1444,8 +1458,10 @@ rspamd_upstream_get_hashed (struct upstream_list *ups, const guint8 *key, guint
static struct upstream*
rspamd_upstream_get_common (struct upstream_list *ups,
- enum rspamd_upstream_rotation default_type,
- const guchar *key, gsize keylen, gboolean forced)
+ struct upstream* except,
+ enum rspamd_upstream_rotation default_type,
+ const guchar *key, gsize keylen,
+ gboolean forced)
{
enum rspamd_upstream_rotation type;
struct upstream *up = NULL;
@@ -1459,6 +1475,12 @@ rspamd_upstream_get_common (struct upstream_list *ups,
}
RSPAMD_UPSTREAM_UNLOCK (ups);
+ if (ups->alive->len == 1) {
+ /* Fast path */
+ up = g_ptr_array_index (ups->alive, 0);
+ goto end;
+ }
+
if (!forced) {
type = ups->rot_alg != RSPAMD_UPSTREAM_UNDEF ? ups->rot_alg : default_type;
}
@@ -1474,16 +1496,16 @@ rspamd_upstream_get_common (struct upstream_list *ups,
switch (type) {
default:
case RSPAMD_UPSTREAM_RANDOM:
- up = rspamd_upstream_get_random (ups);
+ up = rspamd_upstream_get_random (ups, except);
break;
case RSPAMD_UPSTREAM_HASHED:
- up = rspamd_upstream_get_hashed (ups, key, keylen);
+ up = rspamd_upstream_get_hashed (ups, except, key, keylen);
break;
case RSPAMD_UPSTREAM_ROUND_ROBIN:
- up = rspamd_upstream_get_round_robin (ups, TRUE);
+ up = rspamd_upstream_get_round_robin (ups, except, TRUE);
break;
case RSPAMD_UPSTREAM_MASTER_SLAVE:
- up = rspamd_upstream_get_round_robin (ups, FALSE);
+ up = rspamd_upstream_get_round_robin (ups, except, FALSE);
break;
case RSPAMD_UPSTREAM_SEQUENTIAL:
if (ups->cur_elt >= ups->alive->len) {
@@ -1495,6 +1517,7 @@ rspamd_upstream_get_common (struct upstream_list *ups,
break;
}
+end:
if (up) {
up->checked ++;
}
@@ -1507,7 +1530,7 @@ rspamd_upstream_get (struct upstream_list *ups,
enum rspamd_upstream_rotation default_type,
const guchar *key, gsize keylen)
{
- return rspamd_upstream_get_common (ups, default_type, key, keylen, FALSE);
+ return rspamd_upstream_get_common (ups, NULL, default_type, key, keylen, FALSE);
}
struct upstream*
@@ -1515,7 +1538,15 @@ rspamd_upstream_get_forced (struct upstream_list *ups,
enum rspamd_upstream_rotation forced_type,
const guchar *key, gsize keylen)
{
- return rspamd_upstream_get_common (ups, forced_type, key, keylen, TRUE);
+ return rspamd_upstream_get_common (ups, NULL, forced_type, key, keylen, TRUE);
+}
+
+struct upstream *rspamd_upstream_get_except (struct upstream_list *ups,
+ struct upstream *except,
+ enum rspamd_upstream_rotation default_type,
+ const guchar *key, gsize keylen)
+{
+ return rspamd_upstream_get_common (ups, except, default_type, key, keylen, FALSE);
}
void
diff --git a/src/libutil/upstream.h b/src/libutil/upstream.h
index 4dd69e0dd..0581a314c 100644
--- a/src/libutil/upstream.h
+++ b/src/libutil/upstream.h
@@ -283,6 +283,17 @@ struct upstream *rspamd_upstream_get_forced (struct upstream_list *ups,
enum rspamd_upstream_rotation forced_type,
const guchar *key, gsize keylen);
+/**
+ * Get new upstream from the list excepting the upstream specified
+ * @param ups upstream list
+ * @param type type of rotation algorithm, for `RSPAMD_UPSTREAM_HASHED` it is required to specify `key` and `keylen` as arguments
+ * @return
+ */
+struct upstream *rspamd_upstream_get_except (struct upstream_list *ups,
+ struct upstream *except,
+ enum rspamd_upstream_rotation default_type,
+ const guchar *key, gsize keylen);
+
/**
* Re-resolve addresses for all upstreams registered
*/
More information about the Commits
mailing list