Fix the buffering used for parsing. (#50)

fill_buf didn't work like I expected. This code is much better anyway.
This commit is contained in:
kixelated 2023-08-02 11:41:28 -07:00 committed by GitHub
parent 0e239935a6
commit 3a65873055
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 67 additions and 65 deletions

1
Cargo.lock generated
View File

@ -967,6 +967,7 @@ name = "moq-warp"
version = "0.1.0" version = "0.1.0"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"bytes",
"log", "log",
"moq-transport", "moq-transport",
"moq-transport-quinn", "moq-transport-quinn",

View File

@ -1,11 +1,10 @@
use anyhow::Context;
use moq_transport::{Decode, DecodeError, Encode, Message}; use moq_transport::{Decode, DecodeError, Encode, Message};
use bytes::{Buf, BufMut, BytesMut}; use bytes::{Buf, BytesMut};
use std::io::Cursor; use std::io::Cursor;
use std::sync::Arc; use std::sync::Arc;
use tokio::sync::Mutex; use tokio::{io::AsyncReadExt, sync::Mutex};
use webtransport_quinn::{RecvStream, SendStream}; use webtransport_quinn::{RecvStream, SendStream};
@ -86,8 +85,7 @@ impl RecvControl {
} }
Err(DecodeError::UnexpectedEnd) => { Err(DecodeError::UnexpectedEnd) => {
// The decode failed, so we need to append more data. // The decode failed, so we need to append more data.
let chunk = self.stream.read_chunk(1024, true).await?.context("stream closed")?; self.stream.read_buf(&mut self.buf).await?;
self.buf.put(chunk.bytes);
} }
Err(e) => return Err(e.into()), Err(e) => return Err(e.into()),
} }

View File

@ -1,12 +1,12 @@
use std::{collections::BinaryHeap, io::Cursor, sync::Arc}; use std::{collections::BinaryHeap, io::Cursor, sync::Arc};
use anyhow::Context; use anyhow::Context;
use bytes::BytesMut; use bytes::{Buf, BytesMut};
use moq_transport::{Decode, DecodeError, Encode, Object}; use moq_transport::{Decode, DecodeError, Encode, Object};
use tokio::io::AsyncWriteExt; use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::Mutex;
use tokio::task::JoinSet; use tokio::task::JoinSet;
use tokio::{io::AsyncBufReadExt, sync::Mutex};
use webtransport_quinn::Session; use webtransport_quinn::Session;
use crate::{RecvStream, SendStream, SendStreamOrder}; use crate::{RecvStream, SendStream, SendStreamOrder};
@ -83,6 +83,8 @@ impl SendObjectsInner {
header.encode(&mut self.buf).unwrap(); header.encode(&mut self.buf).unwrap();
stream.write_all(&self.buf).await.context("failed to write header")?; stream.write_all(&self.buf).await.context("failed to write header")?;
// log::info!("created stream: {:?}", header);
Ok(stream) Ok(stream)
} }
} }
@ -117,18 +119,15 @@ impl RecvObjects {
} }
} }
async fn read(stream: webtransport_quinn::RecvStream) -> anyhow::Result<(Object, RecvStream)> { async fn read(mut stream: webtransport_quinn::RecvStream) -> anyhow::Result<(Object, RecvStream)> {
let mut stream = RecvStream::new(stream); let mut buf = BytesMut::new();
loop { loop {
// Read more data into the buffer. // Read more data into the buffer.
let data = stream.fill_buf().await?; stream.read_buf(&mut buf).await?;
if data.is_empty() {
anyhow::bail!("stream closed before reading header");
}
// Use a cursor to read the buffer and remember how much we read. // Use a cursor to read the buffer and remember how much we read.
let mut read = Cursor::new(data); let mut read = Cursor::new(&mut buf);
let header = match Object::decode(&mut read) { let header = match Object::decode(&mut read) {
Ok(header) => header, Ok(header) => header,
@ -136,10 +135,14 @@ impl RecvObjects {
Err(err) => return Err(err.into()), Err(err) => return Err(err.into()),
}; };
// We parsed a full header, advance the cursor. // We parsed a full header, advance the buffer.
// The borrow checker requires these on separate lines.
let size = read.position() as usize; let size = read.position() as usize;
stream.consume(size); buf.advance(size);
let buf = buf.freeze();
// log::info!("received stream: {:?}", header);
let stream = RecvStream::new(buf, stream);
return Ok((header, stream)); return Ok((header, stream));
} }

View File

@ -1,12 +1,12 @@
use std::{ use std::{
io, io,
ops::{Deref, DerefMut}, pin::{pin, Pin},
pin::Pin,
sync::{Arc, Mutex, Weak}, sync::{Arc, Mutex, Weak},
task, task::{self, Poll},
}; };
use tokio::io::{AsyncWrite, BufReader}; use bytes::{BufMut, Bytes};
use tokio::io::{AsyncRead, AsyncWrite};
// Ugh, so we need to wrap SendStream with a mutex because we need to be able to call set_priority on it. // Ugh, so we need to wrap SendStream with a mutex because we need to be able to call set_priority on it.
// The problem is that set_priority takes a i32, while send_order is a VarInt // The problem is that set_priority takes a i32, while send_order is a VarInt
@ -83,33 +83,33 @@ impl AsyncWrite for SendStream {
} }
// Unfortunately, we need to wrap RecvStream with a buffer since moq-transport::Coding only supports buffered reads. // Unfortunately, we need to wrap RecvStream with a buffer since moq-transport::Coding only supports buffered reads.
// TODO support unbuffered reads so we only read the MoQ header and then hand off the stream. // We first serve any data in the buffer, then we poll the stream.
// NOTE: We can't use AsyncRead::chain because we need to get the inner stream for stop.
pub struct RecvStream { pub struct RecvStream {
stream: BufReader<webtransport_quinn::RecvStream>, buf: Bytes,
stream: webtransport_quinn::RecvStream,
} }
impl RecvStream { impl RecvStream {
pub(crate) fn new(stream: webtransport_quinn::RecvStream) -> Self { pub(crate) fn new(buf: Bytes, stream: webtransport_quinn::RecvStream) -> Self {
let stream = BufReader::new(stream); Self { buf, stream }
Self { stream }
} }
pub fn stop(self, code: u32) { pub fn stop(&mut self, code: u32) {
self.stream.into_inner().stop(code).ok(); self.stream.stop(code).ok();
} }
} }
impl Deref for RecvStream { impl AsyncRead for RecvStream {
type Target = BufReader<webtransport_quinn::RecvStream>; fn poll_read(
mut self: Pin<&mut Self>,
fn deref(&self) -> &Self::Target { cx: &mut task::Context<'_>,
&self.stream buf: &mut tokio::io::ReadBuf<'_>,
} ) -> Poll<io::Result<()>> {
} if !self.buf.is_empty() {
buf.put(&mut pin!(self).buf);
impl DerefMut for RecvStream { Poll::Ready(Ok(()))
fn deref_mut(&mut self) -> &mut Self::Target { } else {
&mut self.stream Pin::new(&mut self.stream).poll_read(cx, buf)
}
} }
} }

View File

@ -1,15 +1,15 @@
[package] [package]
name = "moq-warp" name = "moq-warp"
description = "Media over QUIC" description = "Media over QUIC"
authors = [ "Luke Curley" ] authors = ["Luke Curley"]
repository = "https://github.com/kixelated/moq-rs" repository = "https://github.com/kixelated/moq-rs"
license = "MIT OR Apache-2.0" license = "MIT OR Apache-2.0"
version = "0.1.0" version = "0.1.0"
edition = "2021" edition = "2021"
keywords = [ "quic", "http3", "webtransport", "media", "live" ] keywords = ["quic", "http3", "webtransport", "media", "live"]
categories = [ "multimedia", "network-programming", "web-programming" ] categories = ["multimedia", "network-programming", "web-programming"]
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@ -21,6 +21,7 @@ moq-transport-quinn = { path = "../moq-transport-quinn" }
tokio = "1.27" tokio = "1.27"
anyhow = "1.0.70" anyhow = "1.0.70"
log = "0.4" # TODO remove log = "0.4" # TODO remove
bytes = "1.4"
# QUIC stuff # QUIC stuff
quinn = "0.10" quinn = "0.10"

View File

@ -1,10 +1,5 @@
use super::watch; use super::watch;
use std::sync::Arc; use bytes::Bytes;
// Use Arc to avoid cloning the data for each subscriber. pub type Publisher = watch::Publisher<Bytes>;
pub type Shared = Arc<Vec<u8>>; pub type Subscriber = watch::Subscriber<Bytes>;
// TODO combine fragments into the same buffer, instead of separate buffers.
pub type Publisher = watch::Publisher<Shared>;
pub type Subscriber = watch::Subscriber<Shared>;

View File

@ -1,5 +1,6 @@
use super::{fragment, watch}; use super::watch;
use bytes::Bytes;
use moq_transport::VarInt; use moq_transport::VarInt;
use std::ops::Deref; use std::ops::Deref;
use std::sync::Arc; use std::sync::Arc;
@ -21,7 +22,7 @@ pub struct Publisher {
pub info: Arc<Info>, pub info: Arc<Info>,
// A list of fragments that make up the segment. // A list of fragments that make up the segment.
pub fragments: watch::Publisher<fragment::Shared>, pub fragments: watch::Publisher<Bytes>,
} }
impl Publisher { impl Publisher {
@ -53,7 +54,7 @@ pub struct Subscriber {
pub info: Arc<Info>, pub info: Arc<Info>,
// A list of fragments that make up the segment. // A list of fragments that make up the segment.
pub fragments: watch::Subscriber<fragment::Shared>, pub fragments: watch::Subscriber<Bytes>,
} }
impl Deref for Subscriber { impl Deref for Subscriber {

View File

@ -2,13 +2,15 @@ use std::collections::HashMap;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time; use std::time;
use tokio::io::AsyncBufReadExt; use tokio::io::AsyncReadExt;
use tokio::sync::mpsc; use tokio::sync::mpsc;
use tokio::task::JoinSet; // lock across await boundaries use tokio::task::JoinSet; // lock across await boundaries
use moq_transport::{Announce, AnnounceError, AnnounceOk, Object, Subscribe, SubscribeError, SubscribeOk, VarInt}; use moq_transport::{Announce, AnnounceError, AnnounceOk, Object, Subscribe, SubscribeError, SubscribeOk, VarInt};
use moq_transport_quinn::{RecvObjects, RecvStream}; use moq_transport_quinn::{RecvObjects, RecvStream};
use bytes::BytesMut;
use anyhow::Context; use anyhow::Context;
use super::{broker, control}; use super::{broker, control};
@ -114,16 +116,17 @@ impl Session {
} }
async fn run_segment(mut segment: segment::Publisher, mut stream: RecvStream) -> anyhow::Result<()> { async fn run_segment(mut segment: segment::Publisher, mut stream: RecvStream) -> anyhow::Result<()> {
let mut buf = BytesMut::new();
loop { loop {
let buf = stream.fill_buf().await?; let size = stream.read_buf(&mut buf).await?;
if buf.is_empty() { if size == 0 {
return Ok(()); return Ok(());
} }
let chunk = buf.to_vec(); // Split off the data we read into the buffer, freezing it so multiple threads can read simitaniously.
stream.consume(chunk.len()); let data = buf.split().freeze();
segment.fragments.push(data);
segment.fragments.push(chunk.into())
} }
} }

View File

@ -168,8 +168,8 @@ impl Session {
let mut stream = objects.open(object).await?; let mut stream = objects.open(object).await?;
// Write each fragment as they are available. // Write each fragment as they are available.
while let Some(fragment) = segment.fragments.next().await { while let Some(mut fragment) = segment.fragments.next().await {
stream.write_all(fragment.as_slice()).await?; stream.write_all_buf(&mut fragment).await?;
} }
// NOTE: stream is automatically closed when dropped // NOTE: stream is automatically closed when dropped