add client stream benchmark

Signed-off-by: Xintao <[email protected]>

Author: hunterlxt

benchmark/src/bench.rs

@@ -7,6 +7,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
 use futures::prelude::*;
+use futures::sink::SinkExt;
 use grpc::{
     self, ClientStreamingSink, DuplexSink, MessageReader, Method, MethodType, RequestStream,
     RpcContext, RpcStatus, RpcStatusCode, ServerStreamingSink, ServiceBuilder, UnarySink,
@@ -56,12 +57,21 @@ impl BenchmarkService for Benchmark {
     fn streaming_from_client(
         &mut self,
         ctx: RpcContext,
-        _: RequestStream<SimpleRequest>,
+        mut stream: RequestStream<SimpleRequest>,
         sink: ClientStreamingSink<SimpleResponse>,
     ) {
-        let f = sink.fail(RpcStatus::new(RpcStatusCode::UNIMPLEMENTED, None));
+        let f = async move {
+            let mut resp: Option<SimpleResponse> = None;
+            while let Some(req) = stream.try_next().await? {
+                if resp.is_none() {
+                    resp = Some(gen_resp(&req));
+                }
+            }
+            sink.success(resp.unwrap()).await?;
+            Ok(())
+        };
         let keep_running = self.keep_running.clone();
-        spawn!(ctx, keep_running, "reporting unimplemented method", f)
+        spawn!(ctx, keep_running, "streaming from client", f)
     }
 
     fn streaming_from_server(

benchmark/src/client.rs

@@ -182,6 +182,7 @@ struct RequestExecutor<B> {
     ctx: ExecutorContext<B>,
     client: Arc<BenchmarkServiceClient>,
     req: SimpleRequest,
+    cfg: ClientConfig,
 }
 
 impl<B: Backoff + Send + 'static> RequestExecutor<B> {
@@ -190,6 +191,7 @@ impl<B: Backoff + Send + 'static> RequestExecutor<B> {
             ctx,
             client: Arc::new(BenchmarkServiceClient::new(channel)),
             req: gen_req(cfg),
+            cfg: cfg.clone(),
         }
     }
 
@@ -254,6 +256,38 @@ impl<B: Backoff + Send + 'static> RequestExecutor<B> {
         };
         spawn!(client, keep_running, "streaming ping pong", f);
     }
+
+    fn execute_stream_from_client(mut self) {
+        let client = self.client.clone();
+        let keep_running = self.ctx.keep_running.clone();
+
+        let mut send_data = vec![];
+        for _ in 0..self.cfg.get_messages_per_stream() {
+            send_data.push(self.req.clone());
+        }
+        let f = async move {
+            loop {
+                let (mut sender, receiver) = self.client.streaming_from_client().unwrap();
+                let latency_timer = Instant::now();
+                let send_stream = futures::stream::iter(send_data.clone());
+                sender
+                    .send_all(&mut send_stream.map(move |item| Ok((item, WriteFlags::default()))))
+                    .await?;
+                sender.close().await?;
+                receiver.await?;
+                self.ctx.observe_latency(latency_timer.elapsed());
+                let mut time = self.ctx.backoff_async();
+                if let Some(t) = &mut time {
+                    t.await;
+                }
+                if !self.ctx.keep_running() {
+                    break;
+                }
+            }
+            Ok(())
+        };
+        spawn!(client, keep_running, "streaming from client", f);
+    }
 }
 
 fn execute<B: Backoff + Send + 'static>(
@@ -272,7 +306,7 @@ fn execute<B: Backoff + Send + 'static>(
         ClientType::ASYNC_CLIENT => match cfg.get_rpc_type() {
             RpcType::UNARY => {
                 if cfg.get_payload_config().has_bytebuf_params() {
-                    panic!("only streaming is supported for generic service.");
+                    panic!("only ping pong streaming is supported for generic service.");
                 }
                 RequestExecutor::new(ctx, ch, cfg).execute_unary_async()
             }
@@ -283,6 +317,12 @@ fn execute<B: Backoff + Send + 'static>(
                     RequestExecutor::new(ctx, ch, cfg).execute_stream_ping_pong()
                 }
             }
+            RpcType::STREAMING_FROM_CLIENT => {
+                if cfg.get_payload_config().has_bytebuf_params() {
+                    panic!("only ping pong streaming is supported for generic service.");
+                }
+                RequestExecutor::new(ctx, ch, cfg).execute_stream_from_client()
+            }
             _ => unimplemented!(),
         },
         _ => unimplemented!(),