From: Lukáš Ondráček <lukas.ondracek@nic.cz>
Date: Wed, 10 Apr 2024 16:17:26 +0000 (+0200)
Subject: rrl: modify KRU api to return maximum final load value
X-Git-Tag: v6.0.9~1^2~58
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=0e205e6105df46cdab4703ef568b25f792d7d1a8;p=thirdparty%2Fknot-resolver.git

rrl: modify KRU api to return maximum final load value
---

diff --git a/daemon/rrl/api.c b/daemon/rrl/api.c
index 884edf976..d9f133681 100644
--- a/daemon/rrl/api.c
+++ b/daemon/rrl/api.c
@@ -153,18 +153,19 @@ bool kr_rrl_request_begin(struct kr_request *req)
 	if (the_rrl) {
 		uint8_t key[16] ALIGNED(16) = {0, };
 		uint8_t limited_prefix;
+		// uint16_t max_final_load = 0;  // TODO use for query ordering and/or soft limit with TC=1
 		if (req->qsource.addr->sa_family == AF_INET6) {
 			struct sockaddr_in6 *ipv6 = (struct sockaddr_in6 *)req->qsource.addr;
 			memcpy(key, &ipv6->sin6_addr, 16);
 
 			limited_prefix = KRU.limited_multi_prefix_or((struct kru *)the_rrl->kru, kr_now(),
-					1, key, RRL_V6_PREFIXES, the_rrl->v6_prices, RRL_V6_PREFIXES_CNT);
+					1, key, RRL_V6_PREFIXES, the_rrl->v6_prices, RRL_V6_PREFIXES_CNT, /* &max_final_load */ NULL);
 		} else {
 			struct sockaddr_in *ipv4 = (struct sockaddr_in *)req->qsource.addr;
 			memcpy(key, &ipv4->sin_addr, 4);  // TODO append port?
 
 			limited_prefix = KRU.limited_multi_prefix_or((struct kru *)the_rrl->kru, kr_now(),
-					0, key, RRL_V4_PREFIXES, the_rrl->v4_prices, RRL_V4_PREFIXES_CNT);
+					0, key, RRL_V4_PREFIXES, the_rrl->v4_prices, RRL_V4_PREFIXES_CNT, /* &max_final_load */ NULL);
 		}
 		limited = limited_prefix;
 	}
diff --git a/daemon/rrl/kru.h b/daemon/rrl/kru.h
index de678d876..6972fe8e1 100644
--- a/daemon/rrl/kru.h
+++ b/daemon/rrl/kru.h
@@ -76,8 +76,9 @@ struct kru_api {
 	/// Updates KRU only if no query is blocked, unless a race condition occurs --
 	/// in such a case all longer prefixes might have been updated.
 	/// The key of i-th query consists of prefixes[i] bits of key, prefixes[i], and namespace.
+	/// If zero is returned, *max_load_out is set to the maximum of final values of the involved counters normalized to the limit 2^16.
 	uint8_t (*limited_multi_prefix_or)(struct kru *kru, uint32_t time_now,
-			uint8_t namespace, uint8_t key[static 16], uint8_t *prefixes, kru_price_t *prices, size_t queries_cnt);
+			uint8_t namespace, uint8_t key[static 16], uint8_t *prefixes, kru_price_t *prices, size_t queries_cnt, uint16_t *max_load_out);
 };
 // The functions are stored this way to make it easier to switch
 // implementation based on detected CPU.
diff --git a/daemon/rrl/kru.inc.c b/daemon/rrl/kru.inc.c
index b3a662af3..bd43cab1e 100644
--- a/daemon/rrl/kru.inc.c
+++ b/daemon/rrl/kru.inc.c
@@ -70,6 +70,7 @@ inline static uint64_t rand_bits(unsigned int bits) {
 
 #include "./kru-decay.inc.c"
 
+#include "contrib/ucw/lib.h"
 #include "libdnssec/error.h"
 #include "libdnssec/random.h"
 typedef uint64_t hash_t;
@@ -167,6 +168,7 @@ struct query_ctx {
 	uint16_t price16, limit16;
 	uint16_t id;
 	uint16_t *load;
+	uint16_t final_load_value;  // set by kru_limited_update if not blocked
 };
 
 /// Phase 1/3 of a query -- hash, prefetch, ctx init. Based on one 16-byte key.
@@ -412,6 +414,8 @@ static inline bool kru_limited_update(struct kru *kru, struct query_ctx *ctx)
 		if (load_orig >= limit)
 			return true;
 	} while (!atomic_compare_exchange_weak_explicit(load_at, &load_orig, load_orig + price, memory_order_relaxed, memory_order_relaxed));
+
+	ctx->final_load_value = load_orig + price;
 	return false;
 }
 
@@ -458,7 +462,7 @@ static bool kru_limited_multi_or_nobreak(struct kru *kru, uint32_t time_now, uin
 }
 
 static uint8_t kru_limited_multi_prefix_or(struct kru *kru, uint32_t time_now, uint8_t namespace,
-                                           uint8_t key[static 16], uint8_t *prefixes, kru_price_t *prices, size_t queries_cnt)
+                                           uint8_t key[static 16], uint8_t *prefixes, kru_price_t *prices, size_t queries_cnt, uint16_t *max_load_out)
 {
 	struct query_ctx ctx[queries_cnt];
 
@@ -476,6 +480,12 @@ static uint8_t kru_limited_multi_prefix_or(struct kru *kru, uint32_t time_now, u
 			return prefixes[i];
 	}
 
+	if (max_load_out) {
+		for (size_t i = 0; i < queries_cnt; i++) {
+			*max_load_out = MAX(*max_load_out, ctx[i].final_load_value);
+		}
+	}
+
 	return 0;
 }