grpc-benchmark: Worker and server implementation

grpc-benchmark/Cargo.toml

@@ -6,14 +6,20 @@ license = "MIT"
 rust-version = { workspace = true }
 
 [dependencies]
+async-stream = "0.3"
 grpc = { path = "../grpc" }
 grpc-protobuf = { path = "../grpc-protobuf" }
 prost = "0.14"
 prost-types = "0.14"
 protobuf = { version = "4.35.1-release" }
-tonic = { path = "../tonic" }
+tokio = { version = "1.37.0", features = ["sync", "macros"] }
+tokio-stream = "0.1.18"
+tonic = { path = "../tonic", features = ["tls-aws-lc"] }
 tonic-prost = { path = "../tonic-prost" }
 
+[target.'cfg(unix)'.dependencies]
+nix = { version = "0.31.3", features = ["resource"] }
+
 [build-dependencies]
 tonic-prost-build = { path = "../tonic-prost-build" }
 grpc-protobuf-build = { path = "../grpc-protobuf-build" }

grpc-benchmark/src/lib.rs

@@ -23,19 +23,23 @@
  */
 
 #[allow(unused)]
-mod generated {
-    pub(crate) mod benchmark_service_grpc {
+pub mod generated {
+    pub mod benchmark_service_grpc {
         grpc::include_proto!("grpc/testing", "benchmark_service");
     }
 
-    pub(crate) mod services {
-        pub(crate) mod grpc {
-            pub(crate) mod core {
+    pub mod services {
+        pub mod grpc {
+            pub mod core {
                 include!(concat!(env!("OUT_DIR"), "/tonic/grpc.core.rs"));
             }
-            pub(crate) mod testing {
+            pub mod testing {
                 include!(concat!(env!("OUT_DIR"), "/tonic/grpc.testing.rs"));
             }
         }
     }
 }
+
+mod rusage;
+mod server;
+pub mod worker;

grpc-benchmark/src/main.rs

@@ -0,0 +1,85 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+ * IN THE SOFTWARE.
+ *
+ */
+
+use std::net::IpAddr;
+use std::net::Ipv4Addr;
+use std::net::SocketAddr;
+use std::time::Duration;
+
+use grpc_benchmark::generated::services::grpc::testing::worker_service_server::WorkerServiceServer;
+use grpc_benchmark::worker::WorkerServer;
+use tokio::sync::mpsc;
+use tokio::time;
+use tonic::transport::Server;
+
+pub async fn run_worker(worker_port: u16) -> Result<(), Box<dyn std::error::Error>> {
+    let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), worker_port);
+    let (tx, mut rx) = mpsc::channel(1);
+
+    let svc = WorkerServiceServer::new(WorkerServer::new(tx));
+
+    Server::builder()
+        .add_service(svc)
+        .serve_with_shutdown(addr, async {
+            rx.recv().await;
+            // Wait for the quit_worker response to be sent.
+            time::sleep(Duration::from_secs(1)).await;
+        })
+        .await?;
+
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // The default Tokio runtime uses 1 thread per logical processor. While the
+    // testing framework supports specifying the thread count in the test config,
+    // the tests that run on k8s use specific machine sizes and don't depend on
+    // the clients/servers to restrict their resource usage. Tokio doesn't
+    // support nested runtimes, so support for per test thread config is not
+    // presently supported.
+
+    let mut driver_port = None;
+
+    // Skip the first argument (the binary name itself).
+    for arg in std::env::args().skip(1) {
+        if let Some(port_str) = arg.strip_prefix("--driver_port=") {
+            driver_port = Some(port_str.parse::<u16>().unwrap_or_else(|_| {
+                eprintln!("Error: --driver_port must be a valid u16 integer.");
+                std::process::exit(1);
+            }));
+        } else {
+            eprintln!("Warning: Unrecognized argument '{}'", arg);
+        }
+    }
+
+    let Some(dp) = driver_port else {
+        eprintln!("Usage: worker --driver_port=<port>");
+        std::process::exit(1);
+    };
+
+    run_worker(dp).await?;
+
+    Ok(())
+}

grpc-benchmark/src/rusage.rs

@@ -0,0 +1,62 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+ * IN THE SOFTWARE.
+ *
+ */
+
+#[derive(Debug)]
+pub(crate) struct Rusage {
+    user_time_ns: i64,
+    system_time_ns: i64,
+}
+
+impl Rusage {
+    #[cfg(unix)]
+    pub(crate) fn now() -> Result<Self, String> {
+        use nix::sys::resource::UsageWho;
+        use nix::sys::resource::getrusage;
+        use nix::sys::time::TimeValLike;
+
+        let usage =
+            getrusage(UsageWho::RUSAGE_SELF).map_err(|e| format!("failed to get rusage: {}", e))?;
+
+        Ok(Rusage {
+            user_time_ns: usage.user_time().num_nanoseconds(),
+            system_time_ns: usage.system_time().num_nanoseconds(),
+        })
+    }
+
+    #[cfg(not(unix))]
+    pub(crate) fn now() -> Result<Rusage, String> {
+        Ok(Rusage {
+            user_time_ns: 0,
+            system_time_ns: 0,
+        })
+    }
+
+    pub(crate) fn user_time_nanos(&self) -> i64 {
+        self.user_time_ns
+    }
+
+    pub(crate) fn system_time_nanos(&self) -> i64 {
+        self.system_time_ns
+    }
+}