rechunk stream

src/proto.rs

@@ -53,7 +53,7 @@ pub const PRIO_SECONDARY: RequestPriority = 0x01;
 
 pub(crate) type RequestID = u32;
 type ChunkLength = u16;
-pub(crate) const MAX_CHUNK_LENGTH: ChunkLength = 0x4000;
+const MAX_CHUNK_LENGTH: ChunkLength = 0x4000;
 const CHUNK_HAS_CONTINUATION: ChunkLength = 0x8000;
 
 struct SendQueueItem {
@@ -77,6 +77,10 @@ enum DataReader {
 	Streaming {
 		#[pin]
 		reader: AssociatedStream,
+		packet: Vec<u8>,
+		pos: usize,
+		buf: Vec<u8>,
+		eos: bool,
 	},
 }
 
@@ -84,7 +88,13 @@ impl From<Data> for DataReader {
 	fn from(data: Data) -> DataReader {
 		match data {
 			Data::Full(data) => DataReader::Full { data, pos: 0 },
-			Data::Streaming(reader) => DataReader::Streaming { reader },
+			Data::Streaming(reader) => DataReader::Streaming {
+				reader,
+				packet: Vec::new(),
+				pos: 0,
+				buf: Vec::with_capacity(MAX_CHUNK_LENGTH as usize),
+				eos: false,
+			},
 		}
 	}
 }
@@ -107,16 +117,43 @@ impl Stream for DataReader {
 					Poll::Ready(Some((body, len)))
 				}
 			}
-			DataReaderProj::Streaming { reader } => {
-				reader.poll_next(cx).map(|opt| {
-					opt.map(|v| {
+			DataReaderProj::Streaming {
+				mut reader,
+				packet,
+				pos,
+				buf,
+				eos,
+			} => {
+				if *eos {
+					// eos was reached at previous call to poll_next, where a partial packet
+					// was returned. Now return None
+					return Poll::Ready(None);
+				}
+				loop {
+					let packet_left = packet.len() - *pos;
+					let buf_left = MAX_CHUNK_LENGTH as usize - buf.len();
+					let to_read = std::cmp::min(buf_left, packet_left);
+					buf.extend_from_slice(&packet[*pos..*pos + to_read]);
+					*pos += to_read;
+					if buf.len() == MAX_CHUNK_LENGTH as usize {
+						// we have a full buf, ready to send
+						break;
+					}
+
+					// we don't have a full buf, packet is empty; try receive more
+					if let Some(p) = futures::ready!(reader.as_mut().poll_next(cx)) {
+						*packet = p;
+						*pos = 0;
+					} else {
+						*eos = true;
+						break;
+					}
+				}
+
 				let mut body = [0; MAX_CHUNK_LENGTH as usize];
-						let len = std::cmp::min(MAX_CHUNK_LENGTH as usize, v.len());
-						// TODO this can throw away long vec, they should be splited instead
-						body[..len].copy_from_slice(&v[..len]);
-						(body, len)
-					})
-				})
+				body[..buf.len()].copy_from_slice(&buf);
+				buf.clear();
+				Poll::Ready(Some((body, MAX_CHUNK_LENGTH as usize)))
 			}
 		}
 	}
@@ -196,10 +233,7 @@ pub(crate) trait SendLoop: Sync {
 					data: data.into(),
 				});
 			} else if let Some(mut item) = sending.pop() {
-				trace!(
-					"send_loop: sending bytes for {}",
-                    item.id,
-				);
+				trace!("send_loop: sending bytes for {}", item.id,);
 
 				let data = futures::select! {
 					data = item.data.next().fuse() => data,
@@ -210,7 +244,6 @@ pub(crate) trait SendLoop: Sync {
 
 						// TODO if every SendQueueItem is waiting on data, use select_all to await
 						// something to do
-						// TODO find some way to not require sending empty last chunk
 					}
 				};
 
@@ -222,7 +255,7 @@ pub(crate) trait SendLoop: Sync {
 					None => &[],
 				};
 
-				if !data.is_empty() {
+				if data.len() == MAX_CHUNK_LENGTH as usize {
 					let size_header =
 						ChunkLength::to_be_bytes(data.len() as u16 | CHUNK_HAS_CONTINUATION);
 					write.write_all(&size_header[..]).await?;
@@ -231,7 +264,6 @@ pub(crate) trait SendLoop: Sync {
 
 					sending.push(item);
 				} else {
-					// this is always zero for now, but may be more when above TODO get fixed
 					let size_header = ChunkLength::to_be_bytes(data.len() as u16);
 					write.write_all(&size_header[..]).await?;
 
@@ -317,10 +349,7 @@ pub(crate) trait RecvLoop: Sync + 'static {
 		R: AsyncReadExt + Unpin + Send + Sync,
 	{
 		let mut receiving: HashMap<RequestID, Vec<u8>> = HashMap::new();
-		let mut streams: HashMap<
-			RequestID,
-            ChannelPair,
-		> = HashMap::new();
+		let mut streams: HashMap<RequestID, ChannelPair> = HashMap::new();
 		loop {
 			trace!("recv_loop: reading packet");
 			let mut header_id = [0u8; RequestID::BITS as usize / 8];