[Feature] Allow forced and non-forced upstreams selection

diff --git a/src/libutil/upstream.c b/src/libutil/upstream.c
index 9187ba9b9945565fb8c2b8e05f6918db6cdda74f..f2209228919a16dceddbc76a50c72d5e8cfa9c26 100644 (file)
--- a/src/libutil/upstream.c
+++ b/src/libutil/upstream.c
@@ -786,10 +786,10 @@ rspamd_upstream_get_hashed (struct upstream_list *ups, const guint8 *key, guint
        return g_ptr_array_index (ups->alive, idx);
 }
 
-struct upstream*
-rspamd_upstream_get (struct upstream_list *ups,
+static struct upstream*
+rspamd_upstream_get_common (struct upstream_list *ups,
                enum rspamd_upstream_rotation default_type,
-               const guchar *key, gsize keylen)
+               const guchar *key, gsize keylen, gboolean forced)
 {
        enum rspamd_upstream_rotation type;
 
@@ -800,7 +800,12 @@ rspamd_upstream_get (struct upstream_list *ups,
        }
        rspamd_mutex_unlock (ups->lock);
 
-       type = ups->rot_alg != RSPAMD_UPSTREAM_UNDEF ? ups->rot_alg : default_type;
+       if (!forced) {
+               type = ups->rot_alg != RSPAMD_UPSTREAM_UNDEF ? ups->rot_alg : default_type;
+       }
+       else {
+               type = default_type != RSPAMD_UPSTREAM_UNDEF ? default_type : ups->rot_alg;
+       }
 
        if (type == RSPAMD_UPSTREAM_HASHED && (keylen == 0 || key == NULL)) {
                /* Cannot use hashed rotation when no key is specified, switch to random */
@@ -830,6 +835,22 @@ rspamd_upstream_get (struct upstream_list *ups,
        return NULL;
 }
 
+struct upstream*
+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);
+}
+
+struct upstream*
+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);
+}
+
 void
 rspamd_upstream_reresolve (struct upstream_ctx *ctx)
 {

diff --git a/src/libutil/upstream.h b/src/libutil/upstream.h
index c3cbf685ad6a0e27447305e7fda8889b2743bbc6..07b9302e8baa25c21011ad57ba5d97bcb8cf7844 100644 (file)
--- a/src/libutil/upstream.h
+++ b/src/libutil/upstream.h
@@ -190,6 +190,16 @@ struct upstream* rspamd_upstream_get (struct upstream_list *ups,
                enum rspamd_upstream_rotation default_type,
                const guchar *key, gsize keylen);
 
+/**
+ * Get new upstream from the list
+ * @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_forced (struct upstream_list *ups,
+               enum rspamd_upstream_rotation forced_type,
+               const guchar *key, gsize keylen);
+
 /**
  * Re-resolve addresses for all upstreams registered
  */