feat: configure subscribers_per_topic for momento pubsub workflow (#363)

Problem
A momento client can only support a maximum of 100 subscriptions due to the server side GRPC restriction on concurrent streams per connection. The current implementation assumes that all subscribers subscribe to all topics.
We are hence unable to test workflows like 10 topics with 100 subscribers each and we also run into issues if the number of subscriptions created is more than 100 * number of clients.

Solution
Add a configuration option for subscribers per topic and add runtime checks to fail if the configuration tries to create more subscriptions than what the clients can support.

Author: krispraws

src/clients/pubsub/momento.rs

@@ -34,11 +34,21 @@ pub fn launch_subscribers(
         if let Component::Topics(topics) = component {
             let poolsize = topics.subscriber_poolsize();
             let concurrency = topics.subscriber_concurrency();
-
+            if concurrency > 100 {
+                eprintln!("Momento sdk does not support concurrency values greater than 100.");
+                std::process::exit(1);
+            }
+            let num_topics = topics.topics().len();
+            let subscribers_per_topic = topics
+                .momento_subscribers_per_topic()
+                .unwrap_or(poolsize * concurrency);
+            if num_topics * subscribers_per_topic > poolsize * concurrency {
+                eprintln!("Not enough Momento clients to support the workload - adjust momento_subscribers_per_topic or increase subscriber_poolsize/subscriber_concurrency.");
+                std::process::exit(1);
+            }
+            let mut clients = Vec::<Arc<TopicClient>>::with_capacity(poolsize);
             for _ in 0..poolsize {
                 let client = {
-                    let _guard = runtime.enter();
-
                     // initialize the Momento topic client
                     if std::env::var("MOMENTO_API_KEY").is_err() {
                         eprintln!("environment variable `MOMENTO_API_KEY` is not set");
@@ -53,7 +63,7 @@ pub fn launch_subscribers(
                                 std::process::exit(1);
                             }
                         };
-
+                    let _guard = runtime.enter();
                     match TopicClient::builder()
                         .configuration(LowLatency::v1())
                         .credential_provider(credential_provider)
@@ -66,15 +76,20 @@ pub fn launch_subscribers(
                         }
                     }
                 };
-
-                for _ in 0..concurrency {
-                    for topic in topics.topics() {
-                        runtime.spawn(subscriber_task(
-                            client.clone(),
-                            cache_name.clone(),
-                            topic.to_string(),
-                        ));
-                    }
+                clients.push(client);
+            }
+            let mut client_index = 0;
+            for topic in topics.topics() {
+                for _ in 0..subscribers_per_topic {
+                    // Round-robin over the clients to pick one
+                    let client = &clients[client_index];
+                    client_index = (client_index + 1) % clients.len();
+                    let _guard = runtime.enter();
+                    runtime.spawn(subscriber_task(
+                        client.clone(),
+                        cache_name.clone(),
+                        topic.to_string(),
+                    ));
                 }
             }
         }

src/config/workload.rs

@@ -168,6 +168,8 @@ pub struct Topics {
     topic_distribution: Distribution,
     #[serde(default)]
     kafka_single_subscriber_group: bool,
+    #[serde(default)]
+    momento_subscribers_per_topic: Option<usize>,
 }
 
 impl Topics {
@@ -222,6 +224,10 @@ impl Topics {
     pub fn kafka_single_subscriber_group(&self) -> bool {
         self.kafka_single_subscriber_group
     }
+
+    pub fn momento_subscribers_per_topic(&self) -> Option<usize> {
+        self.momento_subscribers_per_topic
+    }
 }
 
 #[derive(Clone, Deserialize)]

src/workload/mod.rs

@@ -579,6 +579,7 @@ pub struct Topics {
     subscriber_poolsize: usize,
     subscriber_concurrency: usize,
     kafka_single_subscriber_group: bool,
+    momento_subscribers_per_topic: Option<usize>,
 }
 
 impl Topics {
@@ -643,6 +644,7 @@ impl Topics {
             subscriber_poolsize,
             subscriber_concurrency,
             kafka_single_subscriber_group: topics.kafka_single_subscriber_group(),
+            momento_subscribers_per_topic: topics.momento_subscribers_per_topic(),
         }
     }
 
@@ -669,6 +671,10 @@ impl Topics {
     pub fn kafka_single_subscriber_group(&self) -> bool {
         self.kafka_single_subscriber_group
     }
+
+    pub fn momento_subscribers_per_topic(&self) -> Option<usize> {
+        self.momento_subscribers_per_topic
+    }
 }
 
 #[derive(Clone)]