From b914d2f1ea7cc9ebbc5ebf19ad6c0909a281c343 Mon Sep 17 00:00:00 2001
From: Yazad Daruvala <y.daruvala@gmail.com>
Date: Sat, 17 Jun 2017 19:14:30 -0700
Subject: [PATCH] test(server): add benches for TransferEncoding::chunked

---
 benches/server.rs | 108 +++++++++++++++++++++++++++++++++-------------
 1 file changed, 79 insertions(+), 29 deletions(-)

diff --git a/benches/server.rs b/benches/server.rs
index 8148b67a..38a6ea59 100644
--- a/benches/server.rs
+++ b/benches/server.rs
@@ -7,43 +7,71 @@ extern crate pretty_env_logger;
 extern crate test;
 
 use std::io::{Read, Write};
-use std::net::{TcpListener, TcpStream, SocketAddr};
+use std::net::{TcpListener, TcpStream};
 use std::sync::mpsc;
 
 use futures::Future;
 use futures::sync::oneshot;
 
-use hyper::header::{ContentLength, ContentType};
+use hyper::header::{ContentLength, ContentType, TransferEncoding};
 use hyper::server::{self, Service};
 
+macro_rules! bench_server {
+    ($b:ident, $handler:ident, $total_bytes:expr) => {
+        let (_until_tx, until_rx) = oneshot::channel();
+        let addr = {
+            let (addr_tx, addr_rx) = mpsc::channel();
+            ::std::thread::spawn(move || {
+                let addr = "127.0.0.1:0".parse().unwrap();
+                let srv = hyper::server::Http::new().bind(&addr, || Ok($handler)).unwrap();
+                let addr = srv.local_addr().unwrap();
+                addr_tx.send(addr).unwrap();
+                srv.run_until(until_rx.map_err(|_| ())).unwrap();
+            });
 
+            addr_rx.recv().unwrap()
+        };
 
-#[bench]
-fn bench_server_tcp_throughput(b: &mut test::Bencher) {
-    let (_until_tx, until_rx) = oneshot::channel();
-    let addr = spawn_hello(until_rx);
+        let mut tcp = TcpStream::connect(addr).unwrap();
+        let mut buf = [0u8; 8192];
 
-    let mut tcp = TcpStream::connect(addr).unwrap();
-    let mut buf = [0u8; 4096];
-
-    b.bytes = 130 + 35;
-    b.iter(|| {
-        tcp.write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n").unwrap();
-        let n = tcp.read(&mut buf).unwrap();
-        assert_eq!(n, 130);
-    })
+        $b.bytes = 35 + $total_bytes;
+        $b.iter(|| {
+            tcp.write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n").unwrap();
+            let mut sum = 0;
+            while sum < $total_bytes {
+                sum += tcp.read(&mut buf).unwrap();
+            }
+            assert_eq!(sum, $total_bytes);
+        })
+    }
 }
 
 
 #[bench]
-fn bench_raw_tcp_throughput(b: &mut test::Bencher) {
+fn bench_server_tcp_throughput_fixedsize_small_payload(b: &mut test::Bencher) {
+    bench_server!(b, FixedSizeSmallPayload, 130);
+}
+
+#[bench]
+fn bench_server_tcp_throughput_chunked_small_payload(b: &mut test::Bencher) {
+    bench_server!(b, ChunkedSmallPayload, 148);
+}
+
+#[bench]
+fn bench_server_tcp_throughput_chunked_large_payload(b: &mut test::Bencher) {
+    bench_server!(b, ChunkedLargePayload, 13_000_140);
+}
+
+#[bench]
+fn bench_raw_tcp_throughput_small_payload(b: &mut test::Bencher) {
     let (tx, rx) = mpsc::channel();
     let listener = TcpListener::bind("127.0.0.1:0").unwrap();
     let addr = listener.local_addr().unwrap();
     ::std::thread::spawn(move || {
         let mut sock = listener.accept().unwrap().0;
 
-        let mut buf = [0u8; 4096];
+        let mut buf = [0u8; 8192];
         while rx.try_recv().is_err() {
             sock.read(&mut buf).unwrap();
             sock.write_all(b"\
@@ -69,9 +97,9 @@ fn bench_raw_tcp_throughput(b: &mut test::Bencher) {
     tx.send(()).unwrap();
 }
 
-struct Hello;
+struct FixedSizeSmallPayload;
 
-impl Service for Hello {
+impl Service for FixedSizeSmallPayload {
     type Request = server::Request;
     type Response = server::Response;
     type Error = hyper::Error;
@@ -86,15 +114,37 @@ impl Service for Hello {
     }
 }
 
-fn spawn_hello(until: oneshot::Receiver<()>) -> SocketAddr {
-    let (addr_tx, addr_rx) = mpsc::channel();
-    ::std::thread::spawn(move || {
-        let addr = "127.0.0.1:0".parse().unwrap();
-        let srv = hyper::server::Http::new().bind(&addr, || Ok(Hello)).unwrap();
-        let addr = srv.local_addr().unwrap();
-        addr_tx.send(addr).unwrap();
-        srv.run_until(until.map_err(|_| ())).unwrap();
-    });
+struct ChunkedSmallPayload;
 
-    addr_rx.recv().unwrap()
+impl Service for ChunkedSmallPayload {
+    type Request = server::Request;
+    type Response = server::Response;
+    type Error = hyper::Error;
+    type Future = ::futures::Finished<Self::Response, hyper::Error>;
+    fn call(&self, _req: Self::Request) -> Self::Future {
+        ::futures::finished(
+            server::Response::new()
+                .with_header(TransferEncoding::chunked())
+                .with_header(ContentType::plaintext())
+                .with_body("Hello, World!")
+        )
+    }
 }
+
+struct ChunkedLargePayload;
+
+impl Service for ChunkedLargePayload {
+    type Request = server::Request;
+    type Response = server::Response;
+    type Error = hyper::Error;
+    type Future = ::futures::Finished<Self::Response, hyper::Error>;
+    fn call(&self, _req: Self::Request) -> Self::Future {
+        ::futures::finished(
+            server::Response::new()
+                .with_header(TransferEncoding::chunked())
+                .with_header(ContentType::plaintext())
+                .with_body((0..1_000_000).flat_map(|_| "Hello, World!".bytes()).collect::<Vec<u8>>())
+        )
+    }
+}
+