lib/kru: add optional arg prefix_out for logging

2 files changed, 15 insertions, 5 deletions

diff --git a/lib/kru.h b/lib/kru.h
index 0a3a8e80..b36feeb5 100644
--- a/lib/kru.h
+++ b/lib/kru.h
@@ -70,16 +70,18 @@ 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.
+	/// If zero is returned, *max_load_out (unless NULL) 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, uint16_t *max_load_out);
 
 	/// Multiple queries based on different prefixes of a single key.
-	/// Returns the maximum of final values of the involved counters normalized to the limit 2^16.
+	/// Returns the maximum of final values of the involved counters normalized to the limit 2^16
+	/// and stores the corresponding prefix (value in prefixes) to *prefix_out (unless NULL).
 	/// Set prices to NULL to skip updating; otherwise, KRU is always updated, using maximal allowed value on overflow.
 	/// The key of i-th query consists of prefixes[i] bits of key, prefixes[i], and namespace.
 	uint16_t (*load_multi_prefix_max)(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, uint8_t *prefix_out);
 };
 
 // The functions are stored this way to make it easier to switch
diff --git a/lib/kru.inc.c b/lib/kru.inc.c
index b67d3237..53a9f97b 100644
--- a/lib/kru.inc.c
+++ b/lib/kru.inc.c
@@ -565,7 +565,7 @@ static uint8_t kru_limited_multi_prefix_or(struct kru *kru, uint32_t time_now, u
 }
 
 static uint16_t kru_load_multi_prefix_max(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, uint8_t *prefix_out)
 {
 	struct query_ctx ctx[queries_cnt];
 
@@ -583,10 +583,18 @@ static uint16_t kru_load_multi_prefix_max(struct kru *kru, uint32_t time_now, ui
 		}
 	}
 
+	uint8_t prefix = 0;
 	uint16_t max_load = 0;
 	for (size_t i = 0; i < queries_cnt; i++) {
-		max_load = MAX(max_load, ctx[i].final_load_value);
+		if (max_load < ctx[i].final_load_value) {
+			max_load = ctx[i].final_load_value;
+			prefix = prefixes[i];
+		}
+	}
+	if (prefix_out) {
+		*prefix_out = prefix;
 	}
+
 	return max_load;
 }