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 #[derive(Debug)] 91 pub struct MockStream(pub tokio::io::DuplexStream); 92 93 impl Connected for MockStream { 94 type ConnectInfo = (); 95 96 fn connect_info(&self) -> Self::ConnectInfo {} 97 } 98 99 impl AsyncRead for MockStream { 100 fn poll_read( 101 mut self: Pin<&mut Self>, 102 cx: &mut Context<'_>, 103 buf: &mut ReadBuf<'_>, 104 ) -> Poll<std::io::Result<()>> { 105 Pin::new(&mut self.0).poll_read(cx, buf) 106 } 107 } 108 109 impl AsyncWrite for MockStream { 110 fn poll_write( 111 mut self: Pin<&mut Self>, 112 cx: &mut Context<'_>, 113 buf: &[u8], 114 ) -> Poll<std::io::Result<usize>> { 115 Pin::new(&mut self.0).poll_write(cx, buf) 116 } 117 118 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> { 119 Pin::new(&mut self.0).poll_flush(cx) 120 } 121 122 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> { 123 Pin::new(&mut self.0).poll_shutdown(cx) 124 } 125 } 126 127 #[allow(dead_code)] 128 pub async fn mock_io_channel(client: tokio::io::DuplexStream) -> Channel { 129 let mut client = Some(client); 130 131 Endpoint::try_from("http://[::]:50051") 132 .unwrap() 133 .connect_with_connector(service_fn(move |_: Uri| { 134 let client = client.take().unwrap(); 135 async move { Ok::<_, std::io::Error>(MockStream(client)) } 136 })) 137 .await 138 .unwrap() 139 } 140