src/bin/h2client.rs

   1 use std::future::Future;
   2 use std::pin::Pin;
   3 use std::task::{Context, Poll};
   4
   5 use failure::*;
   6 use futures::future::TryFutureExt;
   7 use futures::stream::Stream;
   8 use tokio::net::TcpStream;
   9
  10 // Simple H2 client to test H2 download speed using h2server.rs
  11
  12 struct Process {
  13     body: h2::RecvStream,
  14     trailers: bool,
  15     bytes: usize,
  16 }
  17
  18 impl Future for Process {
  19     type Output = Result<usize, Error>;
  20
  21     fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
  22         let this = self.get_mut();
  23
  24         loop {
  25             if this.trailers {
  26                 match futures::ready!(this.body.poll_trailers(cx)) {
  27                     Ok(Some(trailers)) => println!("trailers: {:?}", trailers),
  28                     Ok(None) => (),
  29                     Err(err) => return Poll::Ready(Err(Error::from(err))),
  30                 }
  31
  32                 println!("Received {} bytes", this.bytes);
  33
  34                 return Poll::Ready(Ok(this.bytes));
  35             } else {
  36                 match futures::ready!(Pin::new(&mut this.body).poll_next(cx)) {
  37                     Some(Ok(chunk)) => {
  38                         this.body.flow_control().release_capacity(chunk.len())?;
  39                         this.bytes += chunk.len();
  40                         // println!("GOT FRAME {}", chunk.len());
  41                     },
  42                     Some(Err(err)) => return Poll::Ready(Err(Error::from(err))),
  43                     None => {
  44                         this.trailers = true;
  45                     },
  46                 }
  47             }
  48         }
  49     }
  50 }
  51
  52 fn send_request(
  53     mut client: h2::client::SendRequest<bytes::Bytes>,
  54 ) -> impl Future<Output = Result<usize, Error>> {
  55
  56     println!("sending request");
  57
  58     let request = http::Request::builder()
  59         .uri("http://localhost/")
  60         .body(())
  61         .unwrap();
  62
  63     let (response, _stream) = client.send_request(request, true).unwrap();
  64
  65     response
  66         .map_err(Error::from)
  67         .and_then(|response| {
  68             Process { body: response.into_body(), trailers: false, bytes: 0 }
  69         })
  70 }
  71
  72 #[tokio::main]
  73 async fn main() -> Result<(), Error> {
  74
  75     let start = std::time::SystemTime::now();
  76
  77     let conn = TcpStream::connect(std::net::SocketAddr::from(([127,0,0,1], 8008)))
  78         .await?;
  79
  80     let (client, h2) = h2::client::Builder::new()
  81         .initial_connection_window_size(1024*1024*1024)
  82         .initial_window_size(1024*1024*1024)
  83         .max_frame_size(4*1024*1024)
  84         .handshake(conn)
  85         .await?;
  86
  87     tokio::spawn(async move {
  88         if let Err(err) = h2.await {
  89             println!("GOT ERR={:?}", err);
  90         }
  91     });
  92
  93     let mut bytes = 0;
  94     for _ in 0..2000 {
  95         bytes += send_request(client.clone()).await?;
  96     }
  97
  98     let elapsed = start.elapsed().unwrap();
  99     let elapsed = (elapsed.as_secs() as f64) +
 100         (elapsed.subsec_millis() as f64)/1000.0;
 101
 102     println!("Downloaded {} bytes, {} MB/s", bytes, (bytes as f64)/(elapsed*1024.0*1024.0));
 103
 104     Ok(())
 105 }