Try to wrap stream into SCC to factor the errored and completed checks

disk/src/cache.rs

@@ -185,7 +185,7 @@ pub fn pre_binned_bytes_for_http(
     query: &PreBinnedQuery,
 ) -> Result<PreBinnedValueByteStream, Error> {
     info!("pre_binned_bytes_for_http  {:?}  {:?}", query, node);
-    let ret = PreBinnedValueByteStream::new(
+    let ret = super::cache::pbv::pre_binned_value_byte_stream_new(
         query.patch.clone(),
         query.channel.clone(),
         query.agg_kind.clone(),

disk/src/cache/pbv.rs

@@ -9,6 +9,7 @@ use err::Error;
 use futures_core::Stream;
 use futures_util::{FutureExt, StreamExt};
 use netpod::log::*;
+use netpod::streamext::SCC;
 use netpod::{
     AggKind, BinnedRange, Channel, NanoRange, NodeConfig, PreBinnedPatchCoord, PreBinnedPatchIterator,
     PreBinnedPatchRange,
@@ -17,41 +18,32 @@ use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
-pub struct PreBinnedValueByteStream {
+pub type PreBinnedValueByteStream = SCC<PreBinnedValueByteStreamInner>;
+
+pub struct PreBinnedValueByteStreamInner {
     inp: PreBinnedValueStream,
-    errored: bool,
-    completed: bool,
 }
 
-impl PreBinnedValueByteStream {
-    pub fn new(patch: PreBinnedPatchCoord, channel: Channel, agg_kind: AggKind, node_config: &NodeConfig) -> Self {
-        Self {
-            inp: PreBinnedValueStream::new(patch, channel, agg_kind, node_config),
-            errored: false,
-            completed: false,
-        }
-    }
+pub fn pre_binned_value_byte_stream_new(
+    patch: PreBinnedPatchCoord,
+    channel: Channel,
+    agg_kind: AggKind,
+    node_config: &NodeConfig,
+) -> PreBinnedValueByteStream {
+    let s1 = PreBinnedValueStream::new(patch, channel, agg_kind, node_config);
+    let s2 = PreBinnedValueByteStreamInner { inp: s1 };
+    SCC::new(s2)
 }
 
-impl Stream for PreBinnedValueByteStream {
+impl Stream for PreBinnedValueByteStreamInner {
     type Item = Result<Bytes, Error>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
         use Poll::*;
-        if self.completed {
-            panic!("poll_next on completed");
-        }
-        if self.errored {
-            self.completed = true;
-            return Ready(None);
-        }
         match self.inp.poll_next_unpin(cx) {
             Ready(Some(item)) => match make_frame::<PreBinnedHttpFrame>(&item) {
                 Ok(buf) => Ready(Some(Ok(buf.freeze()))),
-                Err(e) => {
-                    self.errored = true;
-                    Ready(Some(Err(e.into())))
-                }
+                Err(e) => Ready(Some(Err(e.into()))),
             },
             Ready(None) => Ready(None),
             Pending => Pending,

netpod/src/lib.rs

@@ -10,6 +10,8 @@ use timeunits::*;
 #[allow(unused_imports)]
 use tracing::{debug, error, info, trace, warn};
 
+pub mod streamext;
+
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct AggQuerySingleChannel {
     pub channel_config: ChannelConfig,

netpod/src/streamext.rs

@@ -0,0 +1,72 @@
+use err::Error;
+use futures_core::Stream;
+use futures_util::StreamExt;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+pub struct SCC<S>
+where
+    S: Stream,
+{
+    inp: S,
+    errored: bool,
+    completed: bool,
+}
+
+impl<S> SCC<S>
+where
+    S: Stream,
+{
+    pub fn new(inp: S) -> Self {
+        Self {
+            inp,
+            errored: false,
+            completed: false,
+        }
+    }
+}
+
+impl<S, I> Stream for SCC<S>
+where
+    S: Stream<Item = Result<I, Error>> + Unpin,
+{
+    type Item = <S as Stream>::Item;
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
+        use Poll::*;
+        if self.completed {
+            panic!("poll_next on completed");
+        } else if self.errored {
+            self.completed = true;
+            Ready(None)
+        } else {
+            match self.inp.poll_next_unpin(cx) {
+                Ready(Some(Ok(k))) => Ready(Some(Ok(k))),
+                Ready(Some(Err(e))) => {
+                    self.errored = true;
+                    Ready(Some(Err(e)))
+                }
+                Ready(None) => {
+                    self.completed = true;
+                    Ready(None)
+                }
+                Pending => Pending,
+            }
+        }
+    }
+}
+
+pub trait IntoSCC<S>
+where
+    S: Stream,
+{
+    fn into_scc(self) -> SCC<S>;
+}
+
+impl<S> IntoSCC<S> for S
+where
+    S: Stream,
+{
+    fn into_scc(self) -> SCC<S> {
+        SCC::new(self)
+    }
+}

retrieval/src/client.rs

@@ -52,21 +52,21 @@ pub async fn get_binned(
             let g = match item {
                 Ok(frame) => {
                     type ExpectedType = disk::cache::BinnedBytesForHttpStreamFrame;
-                    info!("frame  len {}", frame.buf().len());
+                    let n1 = frame.buf().len();
                     match bincode::deserialize::<ExpectedType>(frame.buf()) {
                         Ok(item) => match item {
                             Ok(item) => {
-                                info!("item: {:?}", item);
+                                info!("len {}  item {:?}", n1, item);
                                 bin_count += 1;
                                 Some(Ok(item))
                             }
                             Err(e) => {
-                                error!("error frame: {:?}", e);
+                                error!("len {}  error frame {:?}", n1, e);
                                 Some(Err(e))
                             }
                         },
                         Err(e) => {
-                            error!("bincode error: {:?}", e);
+                            error!("len {}  bincode error {:?}", n1, e);
                             Some(Err(e.into()))
                         }
                     }