add function to get the current load-balance child

Author: alandekok

doc/antora/modules/reference/pages/unlang/load-balance.adoc

@@ -76,10 +76,11 @@ into the `instantiate` section of `radiusd.conf.  This configuration
 is no longer used, and the `sql_all` definition can just be placed as
 a module definition into the `mods-enabled/` directory.
 
-== Load-Balance Expansions
+== Modules and Load-Balance Expansions
 
-When the `sql_all` module is defined as above, it can also be used as
-in a xref:xlat/index.adoc[dynamic expansion]:
+When the `sql_all` module is defined as above, the server also
+registers an sql xref:xlat/index.adoc[dynamic expansion].  The
+expansion performs load-balancing for the SQL expansion.
 
 .Example of Load-Balance SQL module
 [source,unlang]
@@ -91,5 +92,28 @@ Reply-Message := %sql_all("SELECT message FROM table WHERE name='%{User-Name}'")
 The expansion works exactly like a `load-balance` block.  One of the
 modules is chosen to run the expansion, in load-balance fashion.
 
+== Related Functions
+
+The server also defines two related functions:
+
+=== `%interpreter.load-balance.persist()`
+
+This function persists the current load-balance choice across
+multiple packets/ It is mainly useful for RADIUS proxying, when EAP
+is used.
+
+If the function is called from outside of a `load-balance` or
+`redundant-load-balance` section, then it does nothing.
+
+The function takes no arguments, and returns no value.
+
+=== `%interpreter.load-balance.child()`
+
+This function returns a `uint8` from `0..N`, which indicates which
+child is currently running.
+
+If the function is called from outside of a `load-balance` or
+`redundant-load-balance` section, then it does nothing.
+
 // Copyright (C) 2026 Network RADIUS SAS.  Licenced under CC-by-NC 4.0.
 // This documentation was developed by Network RADIUS SAS.

src/lib/unlang/interpret.c

@@ -2349,6 +2349,39 @@ static xlat_action_t unlang_interpret_xlat(TALLOC_CTX *ctx, fr_dcursor_t *out,
 	return XLAT_ACTION_DONE;
 }
 
+/** Tell the load-balance keyword to persist this child selection
+ *
+ * @ingroup xlat_functions
+ */
+static xlat_action_t unlang_interpret_lb_persist_xlat(UNUSED TALLOC_CTX *ctx, UNUSED fr_dcursor_t *out,
+						      UNUSED xlat_ctx_t const *xctx,
+						      request_t *request, UNUSED fr_value_box_list_t *in)
+{
+	if (unlang_load_balance_persist(request) < 0) return XLAT_ACTION_FAIL;
+
+	return XLAT_ACTION_DONE;
+}
+
+
+/** Tell the load-balance keyword to persist this child selection
+ *
+ * @ingroup xlat_functions
+ */
+static xlat_action_t unlang_interpret_lb_child_xlat(TALLOC_CTX *ctx, fr_dcursor_t *out,
+						    UNUSED xlat_ctx_t const *xctx,
+						    request_t *request, UNUSED fr_value_box_list_t *in)
+{
+	fr_value_box_t		*vb;
+
+	MEM(vb = fr_value_box_alloc(ctx, FR_TYPE_UINT8, NULL));
+
+	vb->vb_uint8 = unlang_load_balance_child(request);
+
+	fr_dcursor_append(out, vb);
+	return XLAT_ACTION_DONE;
+}
+
+
 /** Return the current virtual server for this request
  *
  * @param[in] request	To return virtual server for.
@@ -2497,6 +2530,9 @@ int unlang_interpret_init_global(TALLOC_CTX *ctx)
 	if (unlikely((xlat = xlat_func_register(ctx, "interpreter", unlang_interpret_xlat, FR_TYPE_VOID)) == NULL)) return -1;
 	xlat_func_args_set(xlat, unlang_interpret_xlat_args);
 
+	if (unlikely((xlat = xlat_func_register(ctx, "interpreter.load-balance.persist", unlang_interpret_lb_persist_xlat, FR_TYPE_VOID)) == NULL)) return -1;
+	if (unlikely((xlat = xlat_func_register(ctx, "interpreter.load-balance.child", unlang_interpret_lb_child_xlat, FR_TYPE_UINT8)) == NULL)) return -1;
+
 	if (unlikely((xlat = xlat_func_register(ctx, "cancel", unlang_cancel_xlat, FR_TYPE_VOID)) == NULL)) return -1;
 	xlat_func_args_set(xlat, unlang_cancel_xlat_args);
 

src/lib/unlang/load_balance.c

@@ -57,6 +57,29 @@ int unlang_load_balance_persist(request_t *request)
 	return request_data_add_const(request, frame->instruction, 0, child, true);
 }
 
+/**  Returns the current child of the load balance section
+ *
+ *  If the frame is UNLANG_TYPE_LOAD_BALANCE or
+ *  UNLANG_TYPE_REDUNDANT_LOAD_BALANCE, then return the child number.
+ */
+uint8_t unlang_load_balance_child(request_t *request)
+{
+	unlang_stack_t		*stack = request->stack;
+	unlang_stack_frame_t	*frame = &stack->frame[stack->depth];
+	unlang_frame_state_redundant_t	*redundant;
+
+	if (!frame->prev.frame_load_balance) return 0;
+
+	fr_assert(frame->prev.frame_load_balance < stack->depth);
+
+	frame = &stack->frame[frame->prev.frame_load_balance];
+	redundant = talloc_get_type_abort(frame->state, unlang_frame_state_redundant_t);
+
+	fr_assert(redundant->num <= UINT8_MAX);
+
+	return redundant->num;
+}
+
 #define unlang_redundant_load_balance unlang_load_balance
 
 static unlang_action_t unlang_load_balance_next(unlang_result_t *p_result, request_t *request,
@@ -98,7 +121,12 @@ static unlang_action_t unlang_load_balance_next(unlang_result_t *p_result, reque
 	 *	end, loop around to the next one.
 	 */
 	redundant->child = unlang_list_next(&g->children, redundant->child);
-	if (!redundant->child) redundant->child = unlang_list_head(&g->children);
+	if (!redundant->child) {
+		redundant->child = unlang_list_head(&g->children);
+		redundant->num = 0;
+	} else {
+		redundant->num++;
+	}
 
 	/*
 	 *	We looped back to the start.  Return whatever results we had from the last child.
@@ -161,7 +189,24 @@ static unlang_action_t unlang_load_balance(unlang_result_t *p_result, request_t
 	redundant = talloc_get_type_abort(frame->state, unlang_frame_state_redundant_t);
 
 	redundant->start = request_data_get(request, frame->instruction, 0);
-	if (redundant->start) goto selected_child;
+	if (redundant->start) {
+		uint32_t i;
+
+		/*
+		 *	This loop should be small, typically less than 16 items.
+		 */
+		for (i = 0; i < unlang_list_num_elements(&g->children); i++) {
+			if (gext->children[i] != redundant->start) continue;
+
+			redundant->num = i;
+			RDEBUG3("load-balance starting at child %u", redundant->num);
+			goto selected_child;
+		}
+
+		fr_assert(0);
+
+		goto selected_child;
+	}
 
 	if (gext->vpt) {
 		uint32_t start;
@@ -205,9 +250,10 @@ static unlang_action_t unlang_load_balance(unlang_result_t *p_result, request_t
 		}
 		talloc_free(to_free);
 
-		RDEBUG3("load-balance starting at child %d", (int) start);
+		RDEBUG3("load-balance starting at child %u", start);
 
 		redundant->start = gext->children[start];
+		redundant->num = start;
 
 	} else {
 		uint32_t start, one, two;
@@ -231,8 +277,9 @@ static unlang_action_t unlang_load_balance(unlang_result_t *p_result, request_t
 			start = two;
 		}
 
-		RDEBUG3("load-balance starting at child %d", (int) start);
+		RDEBUG3("load-balance starting at child %u", start);
 		redundant->start = gext->children[start];
+		redundant->num = start;
 	}
 
 selected_child:
@@ -282,6 +329,17 @@ static unlang_t *compile_load_balance_subsection(unlang_t *parent, unlang_compil
 
 	g = unlang_generic_to_group(c);
 
+	/*
+	 *	The various State mangling functions need to limit the number of load-balance sections.
+	 *
+	 *	Plus, it doesn't make a lot of sense to have 256 children of a load-balance section.  Just
+	 *	what the heck are they doing?
+	 */
+	if (unlang_list_num_elements(&g->children) > UINT8_MAX) {
+		cf_log_err(cs, "Too many children for %s section", c->name);
+		return NULL;
+	}
+
 	/*
 	 *	Inside of the "modules" section, it's a virtual module.  The key is the third argument, and
 	 *	the "name2" is the module name, which we ignore here.

src/lib/unlang/load_balance_priv.h

@@ -41,6 +41,7 @@ typedef struct {
 typedef struct {
 	unlang_t 		*child;		//!< the current child we're processing
 	unlang_t		*start;		//!< the starting child
+	uint32_t		num;		//!< the current child number
 	unlang_result_t		result;		//!< for intermediate child results
 } unlang_frame_state_redundant_t;
 
@@ -61,6 +62,7 @@ static inline unlang_group_t *unlang_load_balance_to_group(unlang_load_balance_t
 }
 
 int	unlang_load_balance_persist(request_t *request) CC_HINT(nonnull);
+uint8_t	unlang_load_balance_child(request_t *request) CC_HINT(nonnull);
 
 #ifdef __cplusplus
 }

src/tests/keywords/load-balance-uint8

@@ -15,7 +15,13 @@ uint8 key
 foreach i (%range(50)) {
 key := 0
 load-balance %{key} {
-	ok
+	group {
+		if (%interpreter.load-balance.child() != key) {
+			test_fail
+			return
+		}
+		ok
+	}
 	test_fail
 	test_fail
 	test_fail
@@ -27,7 +33,14 @@ if (!ok) {
 key := 1
 load-balance %{key} {
 	test_fail
-	ok
+	group {
+		if (%interpreter.load-balance.child() != key) {
+			test_fail
+			return
+		}
+		ok
+	}
+
 	test_fail
 	test_fail
 }
@@ -39,7 +52,14 @@ key := 2
 load-balance %{key} {
 	test_fail
 	test_fail
-	ok
+	group {
+		if (%interpreter.load-balance.child() != key) {
+			test_fail
+			return
+		}
+		ok
+	}
+
 	test_fail
 }
 if (!ok) {
@@ -51,7 +71,14 @@ load-balance %{key} {
 	test_fail
 	test_fail
 	test_fail
-	ok
+	group {
+		if (%interpreter.load-balance.child() != key) {
+			test_fail
+			return
+		}
+		ok
+	}
+
 }
 if (!ok) {
 	return