[forge tests] also colocate vfns and vns for the pfn tests (#9137)

### Description

For the PFN tests, also colocate VFNs and VNs, and add the PFNs as a single colocated group.

### Test Plan

Ensure the PFNs have the correct index (larger than the VFNs), via forge.

testsuite/forge-cli/src/main.rs

@@ -1891,7 +1891,7 @@ fn pfn_const_tps(
         .with_initial_validator_count(NonZeroUsize::new(7).unwrap())
         .with_initial_fullnode_count(7)
         .with_emit_job(EmitJobRequest::default().mode(EmitJobMode::ConstTps { tps: 100 }))
-        .add_network_test(PFNPerformance::new(add_cpu_chaos, add_network_emulation))
+        .add_network_test(PFNPerformance::new(7, add_cpu_chaos, add_network_emulation))
         .with_genesis_helm_config_fn(Arc::new(|helm_values| {
             // Require frequent epoch changes
             helm_values["chain"]["epoch_duration_secs"] = 300.into();
@@ -1934,7 +1934,7 @@ fn pfn_performance(
     ForgeConfig::default()
         .with_initial_validator_count(NonZeroUsize::new(7).unwrap())
         .with_initial_fullnode_count(7)
-        .add_network_test(PFNPerformance::new(add_cpu_chaos, add_network_emulation))
+        .add_network_test(PFNPerformance::new(7, add_cpu_chaos, add_network_emulation))
         .with_genesis_helm_config_fn(Arc::new(|helm_values| {
             // Require frequent epoch changes
             helm_values["chain"]["epoch_duration_secs"] = 300.into();

testsuite/forge/src/backend/k8s/fullnode.rs

@@ -340,14 +340,15 @@ pub async fn install_public_fullnode<'a>(
     era: String,
     namespace: String,
     use_port_forward: bool,
+    index: usize,
 ) -> Result<(PeerId, K8sNode)> {
     let default_node_config = get_default_pfn_node_config();
 
     let merged_node_config =
         merge_node_config(default_node_config, serde_yaml::to_value(node_config)?)?;
 
     let node_peer_id = node_config.get_peer_id().unwrap_or_else(PeerId::random);
-    let fullnode_name = format!("fullnode-{}", node_peer_id.short_str());
+    let fullnode_name = format!("public-fullnode-{}-{}", index, node_peer_id.short_str());
 
     // create the NodeConfig configmap
     let fullnode_node_config_config_map_name = format!("{}-config", fullnode_name.clone());
@@ -493,7 +494,7 @@ pub async fn install_public_fullnode<'a>(
             .name
             .context("Fullnode StatefulSet does not have metadata.name")?,
         peer_id: node_peer_id,
-        index: 0,
+        index,
         service_name: full_service_name,
         version: version.clone(),
         namespace,
@@ -721,7 +722,6 @@ mod tests {
         // top level args
         let peer_id = PeerId::random();
         let version = Version::new(0, "banana".to_string());
-        let _fullnode_name = "fullnode-".to_string() + &peer_id.to_string();
 
         // create APIs
         let stateful_set_api = Arc::new(MockStatefulSetApi::from_stateful_set(
@@ -752,6 +752,7 @@ mod tests {
             era,
             namespace,
             false,
+            7,
         )
         .await
         .unwrap();
@@ -760,8 +761,13 @@ mod tests {
         assert_eq!(created_peer_id, peer_id);
         assert_eq!(
             created_node.name,
-            format!("fullnode-{}", &peer_id.short_str())
+            format!(
+                "public-fullnode-{}-{}",
+                created_node.index,
+                &peer_id.short_str()
+            )
         );
         assert!(created_node.name.len() < 64); // This is a k8s limit
+        assert_eq!(created_node.index, 7);
     }
 }

testsuite/forge/src/backend/k8s/swarm.rs

@@ -181,6 +181,7 @@ impl K8sSwarm {
                 .clone(),
             self.kube_namespace.clone(),
             self.use_port_forward,
+            self.fullnodes.len(),
         )
         .await?;
         k8snode.start().await?; // actually start the node. if port-forward is enabled, this is when it gets its ephemeral port

testsuite/testcases/src/public_fullnode_performance.rs

@@ -13,29 +13,33 @@ use aptos_forge::{
 use aptos_logger::info;
 use aptos_sdk::move_types::account_address::AccountAddress;
 use aptos_types::PeerId;
+use itertools::{EitherOrBoth, Itertools};
 use rand::{
     rngs::{OsRng, StdRng},
     seq::SliceRandom,
     Rng, SeedableRng,
 };
+use std::iter::once;
 use tokio::runtime::Runtime;
 
 /// A simple test that adds multiple public fullnodes (PFNs) to the swarm
 /// and submits transactions through them. Network emulation chaos can also
 /// be configured for all nodes in the swarm.
 #[derive(Default)]
 pub struct PFNPerformance {
+    num_pfns: u64,
     add_cpu_chaos: bool,
     add_network_emulation: bool,
     shuffle_rng_seed: [u8; 32],
 }
 
 impl PFNPerformance {
-    pub fn new(add_cpu_chaos: bool, add_network_emulation: bool) -> Self {
+    pub fn new(num_pfns: u64, add_cpu_chaos: bool, add_network_emulation: bool) -> Self {
         // Create a random seed for the shuffle RNG
         let shuffle_rng_seed: [u8; 32] = OsRng.gen();
 
         Self {
+            num_pfns,
             add_cpu_chaos,
             add_network_emulation,
             shuffle_rng_seed,
@@ -56,11 +60,9 @@ impl PFNPerformance {
     /// is added to all validators, VFNs and PFNs in the swarm.
     fn create_network_emulation_chaos(&self, swarm: &mut dyn Swarm) -> SwarmNetEm {
         // Gather and shuffle all peers IDs (so that we get random network emulation)
-        let shuffled_peer_ids = self.gather_and_shuffle_peer_ids(swarm);
+        let shuffled_peer_ids = self.gather_and_shuffle_peer_ids_with_colocation(swarm);
 
         // Create network emulation chaos for the swarm
-        // TODO: VFNs and VNs need to be colocated
-        let shuffled_peer_ids = shuffled_peer_ids.iter().map(|id| vec![*id]).collect();
         create_multi_region_swarm_network_chaos(shuffled_peer_ids, None)
     }
 
@@ -80,6 +82,31 @@ impl PFNPerformance {
 
         all_peer_ids
     }
+
+    /// Gathers and shuffles all peer IDs in the swarm, colocating VFNs with their validator
+    fn gather_and_shuffle_peer_ids_with_colocation(
+        &self,
+        swarm: &mut dyn Swarm,
+    ) -> Vec<Vec<AccountAddress>> {
+        // Identify the validators and fullnodes in the swarm
+        let validator_peer_ids = swarm.validators().map(|v| v.peer_id()).collect::<Vec<_>>();
+        let fullnode_peer_ids = swarm.full_nodes().map(|v| v.peer_id()).collect::<Vec<_>>();
+        let (vfn_peer_ids, pfn_peer_ids) =
+            fullnode_peer_ids.split_at(fullnode_peer_ids.len() - self.num_pfns as usize);
+        let mut vfn_and_vn_ids: Vec<_> = validator_peer_ids
+            .iter()
+            .zip_longest(vfn_peer_ids)
+            .map(|either_or_both| match either_or_both {
+                EitherOrBoth::Both(validator, vfn) => vec![*validator, *vfn],
+                EitherOrBoth::Left(validator) => vec![*validator],
+                EitherOrBoth::Right(_) => panic!("Unexpected"),
+            })
+            .collect();
+        vfn_and_vn_ids.shuffle(&mut StdRng::from_seed(self.shuffle_rng_seed));
+
+        // All PFNs in the first region
+        once(pfn_peer_ids.to_vec()).chain(vfn_and_vn_ids).collect()
+    }
 }
 
 impl Test for PFNPerformance {
@@ -99,8 +126,7 @@ impl NetworkLoadTest for PFNPerformance {
     /// the swarm; and (ii) use those PFNs as the load destination.
     fn setup(&self, ctx: &mut NetworkContext) -> Result<LoadDestination> {
         // Add the PFNs to the swarm
-        let num_pfns = 7;
-        let pfn_peer_ids = create_and_add_pfns(ctx, num_pfns)?;
+        let pfn_peer_ids = create_and_add_pfns(ctx, self.num_pfns)?;
 
         // Add CPU chaos to the swarm
         if self.add_cpu_chaos {