src/bin/h2s-client.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
   9 // Simple H2 client to test H2 download speed using h2s-server.rs
  10
  11 struct Process {
  12     body: h2::RecvStream,
  13     trailers: bool,
  14     bytes: usize,
  15 }
  16
  17 impl Future for Process {
  18     type Output = Result<usize, Error>;
  19
  20     fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
  21         let this = self.get_mut();
  22
  23         loop {
  24             if this.trailers {
  25                 match futures::ready!(this.body.poll_trailers(cx)) {
  26                     Ok(Some(trailers)) => println!("trailers: {:?}", trailers),
  27                     Ok(None) => (),
  28                     Err(err) => return Poll::Ready(Err(Error::from(err))),
  29                 }
  30
  31                 println!("Received {} bytes", this.bytes);
  32
  33                 return Poll::Ready(Ok(this.bytes));
  34             } else {
  35                 match futures::ready!(Pin::new(&mut this.body).poll_next(cx)) {
  36                     Some(Ok(chunk)) => {
  37                         this.body.release_capacity().release_capacity(chunk.len())?;
  38                         this.bytes += chunk.len();
  39                         // println!("GOT FRAME {}", chunk.len());
  40                     },
  41                     Some(Err(err)) => return Poll::Ready(Err(Error::from(err))),
  42                     None => {
  43                         this.trailers = true;
  44                     },
  45                 }
  46             }
  47         }
  48     }
  49 }
  50
  51 fn send_request(
  52     mut client: h2::client::SendRequest<bytes::Bytes>,
  53 ) -> impl Future<Output = Result<usize, Error>> {
  54     println!("sending request");
  55
  56     let request = http::Request::builder()
  57         .uri("http://localhost/")
  58         .body(())
  59         .unwrap();
  60
  61     let (response, _stream) = client.send_request(request, true).unwrap();
  62
  63     response
  64         .map_err(Error::from)
  65         .and_then(|response| {
  66             Process { body: response.into_body(), trailers: false, bytes: 0 }
  67         })
  68 }
  69
  70 #[tokio::main]
  71 async fn main() -> Result<(), Error> {
  72     let start = std::time::SystemTime::now();
  73
  74     let conn = tokio::net::TcpStream::connect(&"127.0.0.1:8008".parse().unwrap()).await?;
  75
  76     conn.set_nodelay(true).unwrap();
  77     conn.set_recv_buffer_size(1024*1024).unwrap();
  78
  79     use openssl::ssl::{SslConnector, SslMethod};
  80
  81     let mut ssl_connector_builder = SslConnector::builder(SslMethod::tls()).unwrap();
  82     ssl_connector_builder.set_verify(openssl::ssl::SslVerifyMode::NONE);
  83     let conn =
  84         tokio_openssl::connect(
  85             ssl_connector_builder.build().configure()?,
  86             "localhost",
  87             conn,
  88         )
  89         .await
  90         .map_err(|err| format_err!("connect failed - {}", err))?;
  91
  92     let (client, h2) = h2::client::Builder::new()
  93         .initial_connection_window_size(1024*1024*1024)
  94         .initial_window_size(1024*1024*1024)
  95         .max_frame_size(4*1024*1024)
  96         .handshake(conn)
  97         .await?;
  98
  99     // Spawn a task to run the conn...
 100     tokio::spawn(async move {
 101         if let Err(e) = h2.await {
 102             println!("GOT ERR={:?}", e);
 103         }
 104     });
 105
 106     let mut bytes = 0;
 107     for _ in 0..100 {
 108         match send_request(client.clone()).await {
 109             Ok(b) => {
 110                 bytes += b;
 111             }
 112             Err(e) => {
 113                 println!("ERROR {}", e);
 114                 return Ok(());
 115             }
 116         }
 117     }
 118
 119     let elapsed = start.elapsed().unwrap();
 120     let elapsed = (elapsed.as_secs() as f64) +
 121         (elapsed.subsec_millis() as f64)/1000.0;
 122
 123     println!("Downloaded {} bytes, {} MB/s", bytes, (bytes as f64)/(elapsed*1024.0*1024.0));
 124
 125     Ok(())
 126 }