replace async_stream in bgpdumper

Author: yuyuyureka

Cargo.lock

@@ -10,7 +10,6 @@ repository = "https://github.com/wobcom/fernglas"
 
 [dependencies]
 anyhow = "1.0"
-async-stream = "0.3"
 async-trait = "0.1"
 axum = { version = "0.7", default-features = false, features = ["query", "http1", "tokio"] }
 bitvec = "1.0"

Cargo.toml

@@ -2,11 +2,12 @@ use crate::bgpdumper::BgpDumper;
 use crate::route_distinguisher::RouteDistinguisher;
 use crate::store::{Client, RouteState, SessionId, Store, TableSelector, TableType};
 use futures_util::future::join_all;
-use futures_util::{pin_mut, StreamExt};
+use futures_util::TryStreamExt;
 use log::*;
 use serde::Deserialize;
 use std::collections::HashMap;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+use std::pin::pin;
 use tokio::net::TcpListener;
 use tokio::net::TcpStream;
 use zettabgp::prelude::BgpNotificationMessage;
@@ -43,8 +44,7 @@ pub async fn run_peer(
         stream,
     );
     let open_message = dumper.start_active().await?;
-    let stream = dumper.lifecycle();
-    pin_mut!(stream);
+    let mut stream = dumper.lifecycle();
     let client_name = cfg
         .name_override
         .or(open_message.caps.iter().find_map(|x| {
@@ -70,11 +70,11 @@ pub async fn run_peer(
         )
         .await;
     loop {
-        let update = match stream.next().await {
-            Some(Ok(update)) => update,
-            Some(Err(Ok(notification))) => break Ok(notification),
-            Some(Err(Err(e))) => anyhow::bail!(e),
-            None => panic!(),
+        let update = match pin!(stream.try_next()).await {
+            Ok(Some(update)) => update,
+            Ok(None) => panic!(),
+            Err(Ok(notification)) => break Ok(notification),
+            Err(Err(e)) => anyhow::bail!(e),
         };
         store
             .insert_bgp_update(

src/bgp_collector.rs

@@ -2,8 +2,9 @@
 // Melnikov, which is licensed under the MIT license.
 
 use bytes::{Buf, BytesMut};
-use futures_util::Stream;
 use futures_util::StreamExt;
+use futures_util::TryStream;
+use futures_util::TryStreamExt;
 use log::*;
 use std::sync::Arc;
 use tokio::io::AsyncWriteExt;
@@ -51,7 +52,11 @@ impl BgpDumper {
             .params
             .prepare_message_buf(&mut buf, BgpMessageType::Open, messagelen)?;
         self.write.lock().await.write_all(&buf[0..blen]).await?;
-        let (msgtype, buf) = self.next_message().await?;
+        let (msgtype, buf) = self
+            .messages()
+            .try_next()
+            .await?
+            .ok_or(BgpError::static_str("Expected open message"))?;
         if msgtype != BgpMessageType::Open {
             return Err(BgpError::static_str("Invalid state to start_active"));
         }
@@ -84,51 +89,61 @@ impl BgpDumper {
         });
         tx
     }
-    async fn next_message(&mut self) -> Result<(BgpMessageType, BytesMut), BgpError> {
-        let mut buf = self
-            .read
-            .next()
-            .await
-            .ok_or(BgpError::static_str("unexpected end of stream"))??;
-        let msg = self.params.decode_message_head(&buf)?;
-        buf.advance(19);
-        buf.truncate(msg.1);
-        Ok((msg.0, buf))
+    fn messages(
+        &mut self,
+    ) -> impl TryStream<Ok = (BgpMessageType, BytesMut), Error = BgpError> + Send + Unpin + use<'_>
+    {
+        (&mut self.read).map(|buf| {
+            let mut buf = buf?;
+            let msg = self.params.decode_message_head(&buf)?;
+            buf.advance(19);
+            buf.truncate(msg.1);
+            Ok((msg.0, buf))
+        })
     }
     pub fn lifecycle(
-        mut self,
-    ) -> impl Stream<Item = Result<BgpUpdateMessage, Result<BgpNotificationMessage, BgpError>>> + Send
-    {
+        &mut self,
+    ) -> impl TryStream<Ok = BgpUpdateMessage, Error = Result<BgpNotificationMessage, BgpError>>
+           + Unpin
+           + Send
+           + use<'_> {
         self.stop_keepalives = Some(self.start_keepalives());
+        let params = self.params.clone();
+        Box::pin(
+            self.messages()
+                .map_err(Err)
+                .try_filter_map(move |(msgtype, buf)| {
+                    let params = params.clone();
+                    async move {
+                        if msgtype == BgpMessageType::Keepalive {
+                            return Ok(None);
+                        }
+                        match msgtype {
+                            BgpMessageType::Open => {
+                                Err(Err(BgpError::static_str("Incorrect open message")))
+                            }
+                            BgpMessageType::Keepalive => Ok(None),
+                            BgpMessageType::Notification => {
+                                let mut msgnotification = BgpNotificationMessage::new();
+                                msgnotification
+                                    .decode_from(&params, &buf[..])
+                                    .map_err(Err)?;
+                                Err(Ok(msgnotification))
+                            }
+                            BgpMessageType::Update => {
+                                let mut msgupdate = BgpUpdateMessage::new();
+                                if let Err(e) = msgupdate.decode_from(&params, &buf[..]) {
+                                    warn!("BGP update decode error: {:?}", e);
+                                    warn!("{:x?}", &buf[..]);
+                                    return Ok(None);
+                                }
 
-        async_stream::try_stream! {
-            loop {
-                let (msgtype, buf) = self.next_message().await.map_err(Err)?;
-                if msgtype == BgpMessageType::Keepalive {
-                    continue;
-                }
-                match msgtype {
-                    BgpMessageType::Open => {
-                        Err(Err(BgpError::static_str("Incorrect open message")))?;
-                    }
-                    BgpMessageType::Keepalive => {}
-                    BgpMessageType::Notification => {
-                        let mut msgnotification = BgpNotificationMessage::new();
-                        msgnotification.decode_from(&self.params, &buf[..]).map_err(Err)?;
-                        Err(Ok(msgnotification))?;
-                    }
-                    BgpMessageType::Update => {
-                        let mut msgupdate = BgpUpdateMessage::new();
-                        if let Err(e) = msgupdate.decode_from(&self.params, &buf[..]) {
-                            warn!("BGP update decode error: {:?}", e);
-                            warn!("{:x?}", &buf[..]);
-                            continue;
+                                Ok(Some(msgupdate))
+                            }
                         }
-                        yield msgupdate;
                     }
-                }
-            }
-        }
+                }),
+        )
     }
 }
 impl Drop for BgpDumper {