@@ -403,6 +403,62 @@ 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 replicas 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 `zero_replicas_counting` is a safety mechanism, which allows the caller to
419+ /// decide 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 , zero_replicas_counting : ZeroReplicasCounting ) -> Option < u32 > {
422+ // An empty role has no fixed replica count when zeros are treated as None.
423+ if zero_replicas_counting == ZeroReplicasCounting :: TreatAsNone
424+ && self . role_groups . is_empty ( )
425+ {
426+ return None ;
427+ }
428+
429+ self . role_groups
430+ . values ( )
431+ . map ( |rg| match rg. replicas {
432+ None => None ,
433+ Some ( 0 ) if zero_replicas_counting == ZeroReplicasCounting :: TreatAsNone => None ,
434+ // The individual replicas are [`u16`]s, so a [`u32`] sum has plenty of space.
435+ Some ( replicas) => Some ( u32:: from ( replicas) ) ,
436+ } )
437+ . sum ( )
438+ }
439+
440+ /// Returns the estimated total number of replicas across all role groups.
441+ ///
442+ /// Unlike [`Self::fixed_replica_count`], this always returns a value: a role group with an unset
443+ /// (i.e. [`None`]) replica count is assumed to run a single replica. Use this when a best-effort
444+ /// estimate is needed even though the exact number of replicas is not hard-coded.
445+ pub fn estimated_replica_count ( & self ) -> u32 {
446+ self . role_groups
447+ . values ( )
448+ . map ( |rg| u32:: from ( rg. replicas . unwrap_or ( 1 ) ) )
449+ . sum ( )
450+ }
451+ }
452+
453+ /// How explicit zero (`0`) replicas on a role group should be counted
454+ #[ derive( Copy , Clone , Debug , PartialEq , Eq ) ]
455+ pub enum ZeroReplicasCounting {
456+ /// Treat them as what they are: `Some(0)`.
457+ TreatAsZero ,
458+ /// Treat them as if the user configured [`None`].
459+ TreatAsNone ,
460+ }
461+
406462impl < Config , ConfigOverrides , RoleConfig >
407463 Role < Config , ConfigOverrides , RoleConfig , JavaCommonConfig >
408464where
@@ -654,4 +710,104 @@ mod tests {
654710 ]
655711 ) ;
656712 }
713+
714+ #[ test]
715+ fn replica_counts_with_all_replicas_set ( ) {
716+ let role = construct_role_with_replicas ( [ Some ( 3 ) , Some ( 2 ) , Some ( 5 ) ] ) ;
717+
718+ assert_eq ! (
719+ role. fixed_replica_count( ZeroReplicasCounting :: TreatAsZero ) ,
720+ Some ( 10 )
721+ ) ;
722+ assert_eq ! (
723+ role. fixed_replica_count( ZeroReplicasCounting :: TreatAsNone ) ,
724+ Some ( 10 )
725+ ) ;
726+ assert_eq ! ( role. estimated_replica_count( ) , 10 ) ;
727+ }
728+
729+ #[ test]
730+ fn replica_counts_with_one_replica_unset ( ) {
731+ let role = construct_role_with_replicas ( [ Some ( 3 ) , None , Some ( 2 ) ] ) ;
732+
733+ assert_eq ! (
734+ role. fixed_replica_count( ZeroReplicasCounting :: TreatAsZero ) ,
735+ None
736+ ) ;
737+ assert_eq ! (
738+ role. fixed_replica_count( ZeroReplicasCounting :: TreatAsNone ) ,
739+ None
740+ ) ;
741+ assert_eq ! ( role. estimated_replica_count( ) , 6 ) ;
742+ }
743+
744+ #[ test]
745+ fn replica_counts_with_a_zero_replica ( ) {
746+ let role = construct_role_with_replicas ( [ Some ( 3 ) , Some ( 0 ) ] ) ;
747+
748+ assert_eq ! (
749+ role. fixed_replica_count( ZeroReplicasCounting :: TreatAsZero ) ,
750+ Some ( 3 )
751+ ) ;
752+ // With treat_zero_as_none the zero turns the whole count into None.
753+ assert_eq ! (
754+ role. fixed_replica_count( ZeroReplicasCounting :: TreatAsNone ) ,
755+ None
756+ ) ;
757+ assert_eq ! ( role. estimated_replica_count( ) , 3 ) ;
758+ }
759+
760+ #[ test]
761+ fn replica_counts_with_a_single_zero_role_group ( ) {
762+ let role = construct_role_with_replicas ( [ Some ( 0 ) ] ) ;
763+
764+ assert_eq ! (
765+ role. fixed_replica_count( ZeroReplicasCounting :: TreatAsZero ) ,
766+ Some ( 0 )
767+ ) ;
768+ assert_eq ! (
769+ role. fixed_replica_count( ZeroReplicasCounting :: TreatAsNone ) ,
770+ None
771+ ) ;
772+ assert_eq ! ( role. estimated_replica_count( ) , 0 ) ;
773+ }
774+
775+ #[ test]
776+ fn replica_counts_without_role_groups ( ) {
777+ let role = construct_role_with_replicas ( vec ! [ ] ) ;
778+
779+ assert_eq ! (
780+ role. fixed_replica_count( ZeroReplicasCounting :: TreatAsZero ) ,
781+ Some ( 0 )
782+ ) ;
783+ assert_eq ! (
784+ role. fixed_replica_count( ZeroReplicasCounting :: TreatAsNone ) ,
785+ None
786+ ) ;
787+ assert_eq ! ( role. estimated_replica_count( ) , 0 ) ;
788+ }
789+
790+ /// Builds a [`Role`] with one role group per passed `replicas` entry, so tests only need to
791+ /// care about the replica counts that [`Role::fixed_replica_count`] operates on.
792+ fn construct_role_with_replicas (
793+ replicas : impl IntoIterator < Item = Option < u16 > > ,
794+ ) -> Role < ( ) , EmptyConfigOverrides , GenericRoleConfig , GenericCommonConfig > {
795+ Role {
796+ config : CommonConfiguration :: default ( ) ,
797+ role_config : GenericRoleConfig :: default ( ) ,
798+ role_groups : replicas
799+ . into_iter ( )
800+ . enumerate ( )
801+ . map ( |( index, replicas) | {
802+ (
803+ format ! ( "role-group-{index}" ) ,
804+ RoleGroup {
805+ config : CommonConfiguration :: default ( ) ,
806+ replicas,
807+ } ,
808+ )
809+ } )
810+ . collect ( ) ,
811+ }
812+ }
657813}
0 commit comments