From 5d6581cfe0647f2ac9d3bbf170ee6ef5ff6cf39e Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@rspamd.com>
Date: Sat, 7 Dec 2024 15:43:55 +0000
Subject: [Feature] Improve address rotation algorithm

We apply the following algorithm:
1) Get the current element and it's AF
2) If the next element has the same AF, then we just move to the next element
3) If the next element has different AF, then we should find the next element with the same AF
4) If we cannot find such element, then we return the next element (switching AF)
---
 src/libutil/upstream.c | 56 +++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 49 insertions(+), 7 deletions(-)

diff --git a/src/libutil/upstream.c b/src/libutil/upstream.c
index 49570dec0..87121f9d3 100644
--- a/src/libutil/upstream.c
+++ b/src/libutil/upstream.c
@@ -1038,16 +1038,58 @@ rspamd_upstream_dtor(struct upstream *up)
 rspamd_inet_addr_t *
 rspamd_upstream_addr_next(struct upstream *up)
 {
-	unsigned int idx, next_idx;
+	unsigned int idx = up->addrs.cur, next_idx = up->addrs.cur, cur_af,
+				 min_errors, min_errors_idx;
 	struct upstream_addr_elt *e1, *e2;
 
-	do {
-		idx = up->addrs.cur;
-		next_idx = (idx + 1) % up->addrs.addr->len;
-		e1 = g_ptr_array_index(up->addrs.addr, idx);
+	/*
+	 * We apply the following algorithm:
+	 * 1) Get the current element and it's AF
+	 * 2) If the next element has the same AF, then we just move to the next element
+	 * 3) If the next element has different AF, then we should find the next element with the same AF
+	 * 4) If we cannot find such element, then we return the next element (switching AF)
+	 */
+
+	e1 = g_ptr_array_index(up->addrs.addr, up->addrs.cur);
+	cur_af = rspamd_inet_address_get_af(e1->addr);
+	min_errors = e1->errors;
+	min_errors_idx = idx;
+
+	for (;;) {
+		unsigned int new_af;
+		next_idx = (next_idx + 1) % up->addrs.addr->len;
 		e2 = g_ptr_array_index(up->addrs.addr, next_idx);
-		up->addrs.cur = next_idx;
-	} while (e2->errors > e1->errors);
+
+		if (e2->errors < min_errors) {
+			min_errors = e2->errors;
+			min_errors_idx = next_idx;
+		}
+
+		if (next_idx == idx) {
+			/* We did a full circle, so we have to select something else */
+			if (e2->errors == 0) {
+				/* No errors on the current address, so we can use it */
+			}
+			else {
+				/* We have some errors, so we had to select the address with the lowest err count */
+				next_idx = min_errors_idx;
+			}
+
+			/* Always stop on full circle */
+			break;
+		}
+
+		new_af = rspamd_inet_address_get_af(e2->addr);
+
+		if (cur_af == new_af && e2->errors <= e1->errors) {
+			/* Same AF */
+			up->addrs.cur = next_idx;
+			return e2->addr;
+		}
+	}
+
+	e2 = g_ptr_array_index(up->addrs.addr, next_idx);
+	up->addrs.cur = next_idx;
 
 	return e2->addr;
 }
-- 
cgit v1.2.3