Move code

crates/streams/src/framed_bytes.rs

@@ -0,0 +1,136 @@
+use bytes::Buf;
+use bytes::BufMut;
+use bytes::Bytes;
+use bytes::BytesMut;
+use err::thiserror;
+use err::ThisError;
+use futures_util::Stream;
+use futures_util::StreamExt;
+use netpod::log::*;
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+
+const FRAME_HEAD_LEN: usize = 16;
+const FRAME_PAYLOAD_MAX: u32 = 1024 * 1024 * 8;
+const BUF_MAX: usize = (FRAME_HEAD_LEN + FRAME_PAYLOAD_MAX as usize) * 2;
+
+#[allow(unused)]
+macro_rules! trace_parse {
+    ($($arg:tt)*) => {
+        if false {
+            trace!($($arg)*);
+        }
+    };
+}
+
+#[derive(Debug, ThisError)]
+pub enum Error {
+    FrameTooLarge,
+    Logic,
+}
+
+pub type BoxedFramedBytesStream = Pin<Box<dyn Stream<Item = Result<Bytes, Error>> + Send>>;
+
+// TODO move this type decl because it is not specific to cbor
+pub type SitemtyFramedBytesStream = Pin<Box<dyn Stream<Item = Result<Bytes, Error>> + Send>>;
+
+pub enum State {
+    Reading,
+    Done,
+}
+
+pub struct FramedBytesStream<S> {
+    inp: S,
+    buf: BytesMut,
+    state: State,
+}
+
+impl<S, E> FramedBytesStream<S>
+where
+    S: Stream<Item = Result<Bytes, E>> + Unpin,
+    E: Into<Error>,
+{
+    pub fn new(inp: S) -> Self {
+        Self {
+            inp,
+            buf: BytesMut::with_capacity(1024 * 256),
+            state: State::Reading,
+        }
+    }
+
+    fn try_parse(&mut self) -> Result<Option<Bytes>, Error> {
+        trace_parse!("try_parse  self.buf.len() {}", self.buf.len());
+        if self.buf.len() < FRAME_HEAD_LEN {
+            return Ok(None);
+        }
+        let n = u32::from_le_bytes(self.buf[..4].try_into().map_err(|_| Error::Logic)?);
+        trace_parse!("try_parse  n {}", n);
+        if n > FRAME_PAYLOAD_MAX {
+            let e = Error::FrameTooLarge;
+            return Err(e);
+        }
+        let frame_len = FRAME_HEAD_LEN + n as usize;
+        trace_parse!("try_parse  frame_len {}", frame_len);
+        assert!(self.buf.len() <= self.buf.capacity());
+        if self.buf.capacity() < frame_len {
+            let add_max = BUF_MAX - self.buf.capacity().min(BUF_MAX);
+            let nadd = ((frame_len.min(FRAME_PAYLOAD_MAX as usize) - self.buf.len()) * 2).min(add_max);
+            self.buf.reserve(nadd);
+        }
+        let adv = (frame_len + 7) / 8 * 8;
+        trace_parse!("try_parse  adv {}", adv);
+        if self.buf.len() < adv {
+            Ok(None)
+        } else {
+            self.buf.advance(FRAME_HEAD_LEN);
+            let buf = self.buf.split_to(n as usize);
+            self.buf.advance(adv - frame_len);
+            Ok(Some(buf.freeze()))
+        }
+    }
+}
+
+impl<S, E> Stream for FramedBytesStream<S>
+where
+    S: Stream<Item = Result<Bytes, E>> + Unpin,
+    E: Into<Error>,
+{
+    type Item = Result<Bytes, Error>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
+        use Poll::*;
+        loop {
+            break match &self.state {
+                State::Reading => match self.try_parse() {
+                    Ok(Some(x)) => Ready(Some(Ok(x))),
+                    Ok(None) => match self.inp.poll_next_unpin(cx) {
+                        Ready(Some(x)) => match x {
+                            Ok(x) => {
+                                self.buf.put_slice(&x);
+                                continue;
+                            }
+                            Err(e) => {
+                                self.state = State::Done;
+                                Ready(Some(Err(e.into())))
+                            }
+                        },
+                        Ready(None) => {
+                            if self.buf.len() > 0 {
+                                warn!("remaining bytes in input buffer, input closed  len {}", self.buf.len());
+                            }
+                            self.state = State::Done;
+                            Ready(None)
+                        }
+                        Pending => Pending,
+                    },
+                    Err(e) => {
+                        self.state = State::Done;
+                        Ready(Some(Err(e)))
+                    }
+                },
+                State::Done => Ready(None),
+            };
+        }
+    }
+}

crates/streams/src/instrument.rs

@@ -0,0 +1,33 @@
+use futures_util::Stream;
+use futures_util::StreamExt;
+use netpod::log::tracing;
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+
+#[pin_project::pin_project]
+pub struct InstrumentStream<S> {
+    #[pin]
+    inp: S,
+    #[pin]
+    span: tracing::Span,
+}
+
+impl<S> InstrumentStream<S> {
+    pub fn new(inp: S, span: tracing::Span) -> Self {
+        Self { inp, span }
+    }
+}
+
+impl<S> Stream for InstrumentStream<S>
+where
+    S: Stream,
+{
+    type Item = <S as Stream>::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+        let _spg = this.span.enter();
+        this.inp.poll_next_unpin(cx)
+    }
+}

crates/streams/src/lib.rs

@@ -4,8 +4,10 @@ pub mod collect;
 pub mod dtflags;
 pub mod filechunkread;
 pub mod firsterr;
+pub mod framed_bytes;
 pub mod frames;
 pub mod generators;
+pub mod instrument;
 pub mod itemclone;
 pub mod json_stream;
 pub mod lenframed;