perf: replace BytesMut with alloy_rlp::encode

paradigmxyz · Mar 8, 2024 · 2eaec96 · 2eaec96
1 parent 91c7dd0
commit 2eaec96
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 64 deletions.
diff --git a/crates/net/discv4/src/config.rs b/crates/net/discv4/src/config.rs
@@ -6,10 +6,7 @@
 use alloy_rlp::Encodable;
 use reth_net_common::ban_list::BanList;
 use reth_net_nat::{NatResolver, ResolveNatInterval};
-use reth_primitives::{
-    bytes::{Bytes, BytesMut},
-    NodeRecord,
-};
+use reth_primitives::{bytes::Bytes, NodeRecord};
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 use std::{
@@ -80,9 +77,7 @@ impl Discv4Config {
 
     /// Add another key value pair to include in the ENR
     pub fn add_eip868_pair(&mut self, key: impl Into<Vec<u8>>, value: impl Encodable) -> &mut Self {
-        let mut buf = BytesMut::new();
-        value.encode(&mut buf);
-        self.add_eip868_rlp_pair(key, buf.freeze())
+        self.add_eip868_rlp_pair(key, Bytes::from(alloy_rlp::encode(&value)))
     }
 
     /// Add another key value pair to include in the ENR
@@ -237,9 +232,7 @@ impl Discv4ConfigBuilder {
 
     /// Add another key value pair to include in the ENR
     pub fn add_eip868_pair(&mut self, key: impl Into<Vec<u8>>, value: impl Encodable) -> &mut Self {
-        let mut buf = BytesMut::new();
-        value.encode(&mut buf);
-        self.add_eip868_rlp_pair(key, buf.freeze())
+        self.add_eip868_rlp_pair(key, Bytes::from(alloy_rlp::encode(&value)))
     }
 
     /// Add another key value pair to include in the ENR

diff --git a/crates/net/discv4/src/lib.rs b/crates/net/discv4/src/lib.rs
@@ -40,10 +40,7 @@ use discv5::{
 use enr::Enr;
 use parking_lot::Mutex;
 use proto::{EnrRequest, EnrResponse, EnrWrapper};
-use reth_primitives::{
-    bytes::{Bytes, BytesMut},
-    hex, ForkId, PeerId, B256,
-};
+use reth_primitives::{bytes::Bytes, hex, ForkId, PeerId, B256};
 use secp256k1::SecretKey;
 use std::{
     cell::RefCell,
@@ -377,9 +374,7 @@ impl Discv4 {
     ///
     /// If the key already exists, this will update it.
     pub fn set_eip868_rlp(&self, key: Vec<u8>, value: impl alloy_rlp::Encodable) {
-        let mut buf = BytesMut::new();
-        value.encode(&mut buf);
-        self.set_eip868_rlp_pair(key, buf.freeze())
+        self.set_eip868_rlp_pair(key, Bytes::from(alloy_rlp::encode(&value)))
     }
 
     #[inline]

diff --git a/crates/net/discv4/src/proto.rs b/crates/net/discv4/src/proto.rs
@@ -565,10 +565,7 @@ mod tests {
                 udp_port: rng.gen(),
             };
 
-            let mut buf = BytesMut::new();
-            msg.encode(&mut buf);
-
-            let decoded = NodeEndpoint::decode(&mut buf.as_ref()).unwrap();
+            let decoded = NodeEndpoint::decode(&mut alloy_rlp::encode(&msg).as_slice()).unwrap();
             assert_eq!(msg, decoded);
         }
     }
@@ -585,10 +582,7 @@ mod tests {
                 udp_port: rng.gen(),
             };
 
-            let mut buf = BytesMut::new();
-            msg.encode(&mut buf);
-
-            let decoded = NodeEndpoint::decode(&mut buf.as_ref()).unwrap();
+            let decoded = NodeEndpoint::decode(&mut alloy_rlp::encode(&msg).as_slice()).unwrap();
             assert_eq!(msg, decoded);
         }
     }
@@ -606,10 +600,7 @@ mod tests {
                 enr_sq: None,
             };
 
-            let mut buf = BytesMut::new();
-            msg.encode(&mut buf);
-
-            let decoded = Ping::decode(&mut buf.as_ref()).unwrap();
+            let decoded = Ping::decode(&mut alloy_rlp::encode(&msg).as_slice()).unwrap();
             assert_eq!(msg, decoded);
         }
     }
@@ -627,10 +618,7 @@ mod tests {
                 enr_sq: Some(rng.gen()),
             };
 
-            let mut buf = BytesMut::new();
-            msg.encode(&mut buf);
-
-            let decoded = Ping::decode(&mut buf.as_ref()).unwrap();
+            let decoded = Ping::decode(&mut alloy_rlp::encode(&msg).as_slice()).unwrap();
             assert_eq!(msg, decoded);
         }
     }
@@ -648,10 +636,7 @@ mod tests {
                 enr_sq: None,
             };
 
-            let mut buf = BytesMut::new();
-            msg.encode(&mut buf);
-
-            let decoded = Pong::decode(&mut buf.as_ref()).unwrap();
+            let decoded = Pong::decode(&mut alloy_rlp::encode(&msg).as_slice()).unwrap();
             assert_eq!(msg, decoded);
         }
     }
@@ -669,10 +654,7 @@ mod tests {
                 enr_sq: Some(rng.gen()),
             };
 
-            let mut buf = BytesMut::new();
-            msg.encode(&mut buf);
-
-            let decoded = Pong::decode(&mut buf.as_ref()).unwrap();
+            let decoded = Pong::decode(&mut alloy_rlp::encode(&msg).as_slice()).unwrap();
             assert_eq!(msg, decoded);
         }
     }

diff --git a/crates/net/eth-wire/src/p2pstream.rs b/crates/net/eth-wire/src/p2pstream.rs
@@ -185,13 +185,14 @@ where
         &mut self,
         reason: DisconnectReason,
     ) -> Result<(), P2PStreamError> {
-        let mut buf = BytesMut::new();
-        P2PMessage::Disconnect(reason).encode(&mut buf);
         trace!(
             %reason,
             "Sending disconnect message during the handshake",
         );
-        self.inner.send(buf.freeze()).await.map_err(P2PStreamError::Io)
+        self.inner
+            .send(Bytes::from(alloy_rlp::encode(P2PMessage::Disconnect(reason))))
+            .await
+            .map_err(P2PStreamError::Io)
     }
 }
 

diff --git a/crates/primitives/src/net.rs b/crates/primitives/src/net.rs
@@ -74,8 +74,7 @@ mod tests {
     };
 
     use super::*;
-    use alloy_rlp::{Decodable, Encodable};
-    use bytes::BytesMut;
+    use alloy_rlp::{Decodable};
     use rand::{thread_rng, Rng, RngCore};
     use reth_rpc_types::PeerId;
 
@@ -126,10 +125,7 @@ mod tests {
                 id: rng.gen(),
             };
 
-            let mut buf = BytesMut::new();
-            record.encode(&mut buf);
-
-            let decoded = NodeRecord::decode(&mut buf.as_ref()).unwrap();
+            let decoded = NodeRecord::decode(&mut alloy_rlp::encode(&record).as_slice()).unwrap();
             assert_eq!(record, decoded);
         }
     }
@@ -147,10 +143,7 @@ mod tests {
                 id: rng.gen(),
             };
 
-            let mut buf = BytesMut::new();
-            record.encode(&mut buf);
-
-            let decoded = NodeRecord::decode(&mut buf.as_ref()).unwrap();
+            let decoded = NodeRecord::decode(&mut alloy_rlp::encode(&record).as_slice()).unwrap();
             assert_eq!(record, decoded);
         }
     }

diff --git a/crates/rpc/rpc-types/src/net.rs b/crates/rpc/rpc-types/src/net.rs
@@ -168,8 +168,7 @@ impl FromStr for NodeRecord {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use alloy_rlp::{Decodable, Encodable};
-    use bytes::BytesMut;
+    use alloy_rlp::{Decodable};
     use rand::{thread_rng, Rng, RngCore};
 
     #[test]
@@ -219,10 +218,7 @@ mod tests {
                 id: rng.gen(),
             };
 
-            let mut buf = BytesMut::new();
-            record.encode(&mut buf);
-
-            let decoded = NodeRecord::decode(&mut buf.as_ref()).unwrap();
+            let decoded = NodeRecord::decode(&mut alloy_rlp::encode(&record).as_slice()).unwrap();
             assert_eq!(record, decoded);
         }
     }
@@ -240,10 +236,7 @@ mod tests {
                 id: rng.gen(),
             };
 
-            let mut buf = BytesMut::new();
-            record.encode(&mut buf);
-
-            let decoded = NodeRecord::decode(&mut buf.as_ref()).unwrap();
+            let decoded = NodeRecord::decode(&mut alloy_rlp::encode(&record).as_slice()).unwrap();
             assert_eq!(record, decoded);
         }
     }