respond to review

Author: suzmue

src/pubsub/benchmarks/throughput/src/main.rs

@@ -43,8 +43,7 @@ async fn main() -> Result<(), anyhow::Error> {
 
 /// Returns true if the benchmark has exceeded its maximum allowed runtime.
 pub(crate) fn done(maximum_runtime: Duration, start: Instant) -> bool {
-    let now = Instant::now();
-    now >= start + maximum_runtime
+    start.elapsed() >= maximum_runtime
 }
 
 /// Returns the current Unix timestamp in milliseconds.

src/pubsub/benchmarks/throughput/src/publisher.rs

@@ -29,10 +29,10 @@ struct Stats {
     send_count: AtomicI64,
     /// Total bytes sent to the client library.
     send_bytes: AtomicI64,
-    /// Total messages successfully acknowledged by the Pub/Sub service.
-    ack_count: AtomicI64,
-    /// Total bytes successfully acknowledged by the Pub/Sub service.
-    ack_bytes: AtomicI64,
+    /// Total messages successfully received by the Pub/Sub service.
+    recv_count: AtomicI64,
+    /// Total bytes successfully received by the Pub/Sub service.
+    recv_bytes: AtomicI64,
     /// Total number of publishing errors encountered.
     error_count: AtomicI64,
 }
@@ -44,29 +44,24 @@ pub async fn run(config: PublisherArgs) -> Result<(), anyhow::Error> {
         config.common.project, config.topic_id
     );
 
-    run_publisher(Arc::new(config.clone()), topic_name.clone()).await;
-
-    Ok(())
-}
-
-/// Creates a Publisher client configured with the benchmark settings.
-async fn create_publisher(config: Arc<PublisherArgs>, topic_name: String) -> Publisher {
-    Publisher::builder(topic_name)
+    let publisher = Publisher::builder(topic_name)
         .set_byte_threshold(config.batch_bytes)
         .set_message_count_threshold(config.batch_size)
         .set_delay_threshold(config.batch_delay)
         .with_grpc_subchannel_count(config.common.grpc_channels)
         .build()
-        .await
-        .unwrap()
+        .await?;
+
+    run_publisher(Arc::new(config.clone()), publisher).await?;
+
+    Ok(())
 }
 
 /// Orchestrates the publishing loop and reporting logic.
 ///
 /// This function spawns a background task to continuously publish messages
 /// and uses the main thread to sleep and report metrics at fixed intervals.
-async fn run_publisher(config: Arc<PublisherArgs>, topic_name: String) {
-    let publisher = create_publisher(config.clone(), topic_name).await;
+async fn run_publisher(config: Arc<PublisherArgs>, publisher: Publisher) -> anyhow::Result<()> {
     let payload_size = config.payload_size;
     let data = bytes::Bytes::from(vec![0u8; payload_size as usize]);
     let semaphore = Arc::new(tokio::sync::Semaphore::new(
@@ -79,26 +74,32 @@ async fn run_publisher(config: Arc<PublisherArgs>, topic_name: String) {
     let publisher_handle = tokio::task::spawn(async move {
         loop {
             // Respect the max_outstanding_messages limit.
-            let permit = semaphore.clone().acquire_owned().await.unwrap();
+            let permit = match semaphore.clone().acquire_owned().await {
+                Ok(p) => p,
+                Err(e) => {
+                    eprintln!("Error acquiring permit: {}", e);
+                    break;
+                }
+            };
             let p = publisher.publish(Message::new().set_data(data.clone()));
             publisher_stats.send_count.fetch_add(1, Ordering::Relaxed);
             publisher_stats
                 .send_bytes
                 .fetch_add(payload_size, Ordering::Relaxed);
 
-            let ack_stats = publisher_stats.clone();
+            let recv_stats = publisher_stats.clone();
             tokio::spawn(async move {
                 let _permit = permit;
                 match p.await {
                     Ok(_) => {
-                        ack_stats.ack_count.fetch_add(1, Ordering::Relaxed);
-                        ack_stats
-                            .ack_bytes
+                        recv_stats.recv_count.fetch_add(1, Ordering::Relaxed);
+                        recv_stats
+                            .recv_bytes
                             .fetch_add(payload_size, Ordering::Relaxed);
                     }
                     Err(e) => {
                         eprintln!("Error: {}", e);
-                        ack_stats.error_count.fetch_add(1, Ordering::Relaxed);
+                        recv_stats.error_count.fetch_add(1, Ordering::Relaxed);
                     }
                 }
             });
@@ -113,26 +114,26 @@ async fn run_publisher(config: Arc<PublisherArgs>, topic_name: String) {
         let timer = Instant::now();
         let start_send_count = stats.send_count.load(Ordering::Relaxed);
         let start_send_bytes = stats.send_bytes.load(Ordering::Relaxed);
-        let start_ack_count = stats.ack_count.load(Ordering::Relaxed);
-        let start_ack_bytes = stats.ack_bytes.load(Ordering::Relaxed);
+        let start_recv_count = stats.recv_count.load(Ordering::Relaxed);
+        let start_recv_bytes = stats.recv_bytes.load(Ordering::Relaxed);
         let start_error_count = stats.error_count.load(Ordering::Relaxed);
 
         tokio::time::sleep(config.common.report_interval).await;
 
         // Calculate deltas since the last report interval.
         let send_count_last = stats.send_count.load(Ordering::Relaxed) - start_send_count;
         let send_bytes_last = stats.send_bytes.load(Ordering::Relaxed) - start_send_bytes;
-        let ack_count_last = stats.ack_count.load(Ordering::Relaxed) - start_ack_count;
-        let ack_bytes_last = stats.ack_bytes.load(Ordering::Relaxed) - start_ack_bytes;
+        let recv_count_last = stats.recv_count.load(Ordering::Relaxed) - start_recv_count;
+        let recv_bytes_last = stats.recv_bytes.load(Ordering::Relaxed) - start_recv_bytes;
         let error_count_last = stats.error_count.load(Ordering::Relaxed) - start_error_count;
         let usage = timer.elapsed();
 
         print_result("Pub", i, send_count_last, send_bytes_last, 0, usage);
         print_result(
-            "Ack",
+            "Recv",
             i,
-            ack_count_last,
-            ack_bytes_last,
+            recv_count_last,
+            recv_bytes_last,
             error_count_last,
             usage,
         );
@@ -141,9 +142,11 @@ async fn run_publisher(config: Arc<PublisherArgs>, topic_name: String) {
     publisher_handle.abort();
 
     println!(
-        "# Publisher: error_count={}, ack_count={}, send_count={}",
+        "# Publisher: error_count={}, received_count={}, send_count={}",
         stats.error_count.load(Ordering::Relaxed),
-        stats.ack_count.load(Ordering::Relaxed),
+        stats.recv_count.load(Ordering::Relaxed),
         stats.send_count.load(Ordering::Relaxed)
     );
+
+    Ok(())
 }

src/pubsub/benchmarks/throughput/src/subscriber.rs

@@ -42,8 +42,7 @@ pub async fn run(config: SubscriberArgs) -> Result<(), anyhow::Error> {
     let subscriber = Subscriber::builder()
         .with_grpc_subchannel_count(config.common.grpc_channels)
         .build()
-        .await
-        .unwrap();
+        .await?;
 
     run_subscriber(Arc::new(config.clone()), subscriber, &subscription_name).await;
 
@@ -91,15 +90,13 @@ async fn run_subscriber(
     let stats = Arc::new(Stats::default());
     let mut tasks = Vec::new();
     let max_outstanding_per_task = config.common.max_outstanding_messages / config.streams;
-    if max_outstanding_per_task > 0 {
-        for _ in 0..config.streams {
-            tasks.push(tokio::spawn(subscriber_task(
-                subscriber.clone(),
-                subscription_name.to_string(),
-                max_outstanding_per_task,
-                stats.clone(),
-            )));
-        }
+    for _ in 0..config.streams {
+        tasks.push(tokio::spawn(subscriber_task(
+            subscriber.clone(),
+            subscription_name.to_string(),
+            max_outstanding_per_task,
+            stats.clone(),
+        )));
     }
 
     let start = Instant::now();