1 use super::*; 2 use bytes::Bytes; 3 use futures::ready; 4 use http_body::Body; 5 use pin_project::pin_project; 6 use std::{ 7 pin::Pin, 8 sync::{ 9 atomic::{AtomicUsize, Ordering::SeqCst}, 10 Arc, 11 }, 12 task::{Context, Poll}, 13 }; 14 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; 15 use tonic::transport::{server::Connected, Channel}; 16 use tower_http::map_request_body::MapRequestBodyLayer; 17 18 /// A body that tracks how many bytes passes through it 19 #[pin_project] 20 pub struct CountBytesBody<B> { 21 #[pin] 22 pub inner: B, 23 pub counter: Arc<AtomicUsize>, 24 } 25 26 impl<B> Body for CountBytesBody<B> 27 where 28 B: Body<Data = Bytes>, 29 { 30 type Data = B::Data; 31 type Error = B::Error; 32 33 fn poll_data( 34 self: Pin<&mut Self>, 35 cx: &mut Context<'_>, 36 ) -> Poll<Option<Result<Self::Data, Self::Error>>> { 37 let this = self.project(); 38 let counter: Arc<AtomicUsize> = this.counter.clone(); 39 match ready!(this.inner.poll_data(cx)) { 40 Some(Ok(chunk)) => { 41 println!("response body chunk size = {}", chunk.len()); 42 counter.fetch_add(chunk.len(), SeqCst); 43 Poll::Ready(Some(Ok(chunk))) 44 } 45 x => Poll::Ready(x), 46 } 47 } 48 49 fn poll_trailers( 50 self: Pin<&mut Self>, 51 cx: &mut Context<'_>, 52 ) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> { 53 self.project().inner.poll_trailers(cx) 54 } 55 56 fn is_end_stream(&self) -> bool { 57 self.inner.is_end_stream() 58 } 59 60 fn size_hint(&self) -> http_body::SizeHint { 61 self.inner.size_hint() 62 } 63 } 64 65 #[allow(dead_code)] 66 pub fn measure_request_body_size_layer( 67 bytes_sent_counter: Arc<AtomicUsize>, 68 ) -> MapRequestBodyLayer<impl Fn(hyper::Body) -> hyper::Body + Clone> { 69 MapRequestBodyLayer::new(move |mut body: hyper::Body| { 70 let (mut tx, new_body) = hyper::Body::channel(); 71 72 let bytes_sent_counter = bytes_sent_counter.clone(); 73 tokio::spawn(async move { 74 while let Some(chunk) = body.data().await { 75 let chunk = chunk.unwrap(); 76 println!("request body chunk size = {}", chunk.len()); 77 bytes_sent_counter.fetch_add(chunk.len(), SeqCst); 78 tx.send_data(chunk).await.unwrap(); 79 } 80 81 if let Some(trailers) = body.trailers().await.unwrap() { 82 tx.send_trailers(trailers).await.unwrap(); 83 } 84 }); 85 86 new_body 87 }) 88 } 89 90 #[allow(dead_code)] 91 pub async fn mock_io_channel(client: tokio::io::DuplexStream) -> Channel { 92 let mut client = Some(client); 93 94 Endpoint::try_from("http://[::]:50051") 95 .unwrap() 96 .connect_with_connector(service_fn(move |_: Uri| { 97 let client = client.take().unwrap(); 98 async move { Ok::<_, std::io::Error>(client) } 99 })) 100 .await 101 .unwrap() 102 } 103