@@ -403,6 +403,39 @@ where
403403 }
404404}
405405
406+ impl < Config , ConfigOverrides , RoleConfig , CommonConfig >
407+ Role < Config , ConfigOverrides , RoleConfig , CommonConfig >
408+ where
409+ RoleConfig : Default + JsonSchema + Serialize ,
410+ CommonConfig : Default + JsonSchema + Serialize ,
411+ ConfigOverrides : Default + JsonSchema + Serialize ,
412+ {
413+ /// Returns [`Some<u32>`] in case the number of nodes is hard-coded to a certain value.
414+ ///
415+ /// This is the case when all `replicas` are set to [`Some<u16>`], in which case they are simply
416+ /// summed.
417+ ///
418+ /// The argument `treat_zero_as_none` is a safety mechanism, which allows the caller to decide
419+ /// if an explicit replica count of `0` should be treated as [`None`]. It also means that
420+ /// [`None`] is returned in case no roleGroups are configured at all.
421+ pub fn fixed_replica_count ( & self , treat_zero_as_none : bool ) -> Option < u32 > {
422+ // An empty role has no fixed replica count when zeros are treated as None.
423+ if treat_zero_as_none && self . role_groups . is_empty ( ) {
424+ return None ;
425+ }
426+
427+ self . role_groups
428+ . values ( )
429+ . map ( |rg| match rg. replicas {
430+ None => None ,
431+ Some ( 0 ) if treat_zero_as_none => None ,
432+ // The individual replicas are [`u16`]s, so a [`u32`] sum has plenty of space.
433+ Some ( replicas) => Some ( u32:: from ( replicas) ) ,
434+ } )
435+ . sum ( )
436+ }
437+ }
438+
406439impl < Config , ConfigOverrides , RoleConfig >
407440 Role < Config , ConfigOverrides , RoleConfig , JavaCommonConfig >
408441where
@@ -654,4 +687,69 @@ mod tests {
654687 ]
655688 ) ;
656689 }
690+
691+ #[ test]
692+ fn fixed_replica_count_sums_all_set_replicas ( ) {
693+ let role = construct_role_with_replicas ( [ Some ( 3 ) , Some ( 2 ) , Some ( 5 ) ] ) ;
694+
695+ assert_eq ! ( role. fixed_replica_count( false ) , Some ( 10 ) ) ;
696+ assert_eq ! ( role. fixed_replica_count( true ) , Some ( 10 ) ) ;
697+ }
698+
699+ #[ test]
700+ fn fixed_replica_count_is_none_if_any_replica_is_unset ( ) {
701+ let role = construct_role_with_replicas ( [ Some ( 3 ) , None , Some ( 2 ) ] ) ;
702+
703+ assert_eq ! ( role. fixed_replica_count( false ) , None ) ;
704+ assert_eq ! ( role. fixed_replica_count( true ) , None ) ;
705+ }
706+
707+ #[ test]
708+ fn fixed_replica_count_treats_zero_according_to_flag ( ) {
709+ let role = construct_role_with_replicas ( [ Some ( 3 ) , Some ( 0 ) ] ) ;
710+
711+ assert_eq ! ( role. fixed_replica_count( false ) , Some ( 3 ) ) ;
712+ // With treat_zero_as_none the zero turns the whole count into None.
713+ assert_eq ! ( role. fixed_replica_count( true ) , None ) ;
714+ }
715+
716+ #[ test]
717+ fn fixed_replica_count_of_single_zero_role_group ( ) {
718+ let role = construct_role_with_replicas ( [ Some ( 0 ) ] ) ;
719+
720+ assert_eq ! ( role. fixed_replica_count( false ) , Some ( 0 ) ) ;
721+ assert_eq ! ( role. fixed_replica_count( true ) , None ) ;
722+ }
723+
724+ #[ test]
725+ fn fixed_replica_count_of_role_without_role_groups_is_zero ( ) {
726+ let role = construct_role_with_replicas ( vec ! [ ] ) ;
727+
728+ assert_eq ! ( role. fixed_replica_count( false ) , Some ( 0 ) ) ;
729+ assert_eq ! ( role. fixed_replica_count( true ) , None ) ;
730+ }
731+
732+ /// Builds a [`Role`] with one role group per passed `replicas` entry, so tests only need to
733+ /// care about the replica counts that [`Role::fixed_replica_count`] operates on.
734+ fn construct_role_with_replicas (
735+ replicas : impl IntoIterator < Item = Option < u16 > > ,
736+ ) -> Role < ( ) , EmptyConfigOverrides , GenericRoleConfig , GenericCommonConfig > {
737+ Role {
738+ config : CommonConfiguration :: default ( ) ,
739+ role_config : GenericRoleConfig :: default ( ) ,
740+ role_groups : replicas
741+ . into_iter ( )
742+ . enumerate ( )
743+ . map ( |( index, replicas) | {
744+ (
745+ format ! ( "role-group-{index}" ) ,
746+ RoleGroup {
747+ config : CommonConfiguration :: default ( ) ,
748+ replicas,
749+ } ,
750+ )
751+ } )
752+ . collect ( ) ,
753+ }
754+ }
657755}
0 commit comments