jleony/nanoreth
1
0
Fork 0
mirror of https://github.com/hl-archive-node/nanoreth.git synced 2025-12-06 10:59:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: remove Bandwidthmeter type (#8698)

Browse Source 
Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
This commit is contained in:
Krishang
2024-06-09 19:10:25 +05:30
committed by GitHub
parent 2852893a7d
commit c7fd507681
9 changed files with 13 additions and 330 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/net/common/Cargo.toml
View File

@ -16,5 +16,4 @@ workspace = true
alloy-primitives.workspace = true
# async
pin-project.workspace = true
tokio = { workspace = true, features = ["full"] }

277
crates/net/common/src/bandwidth_meter.rs
View File

@ -1,277 +0,0 @@
//! Support for metering bandwidth.
//!
//! Takes heavy inspiration from <https://github.com/libp2p/rust-libp2p/blob/master/src/bandwidth.rs>
// Copyright 2019 Parity Technologies (UK) Ltd.
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.
use crate::stream::HasRemoteAddr;
use std::{
convert::TryFrom as _,
io,
net::SocketAddr,
pin::Pin,
sync::{
atomic::{AtomicU64, Ordering},
Arc,
},
task::{ready, Context, Poll},
};
use tokio::{
io::{AsyncRead, AsyncWrite, ReadBuf},
net::TcpStream,
};
/// Meters bandwidth usage of streams
#[derive(Debug)]
struct BandwidthMeterInner {
/// Measures the number of inbound packets
inbound: AtomicU64,
/// Measures the number of outbound packets
outbound: AtomicU64,
}
/// Public shareable struct used for getting bandwidth metering info
#[derive(Clone, Debug)]
pub struct BandwidthMeter {
inner: Arc<BandwidthMeterInner>,
}
impl BandwidthMeter {
/// Returns the total number of bytes that have been downloaded on all the streams.
///
/// > **Note**: This method is by design subject to race conditions. The returned value should
/// > only ever be used for statistics purposes.
pub fn total_inbound(&self) -> u64 {
self.inner.inbound.load(Ordering::Relaxed)
}
/// Returns the total number of bytes that have been uploaded on all the streams.
///
/// > **Note**: This method is by design subject to race conditions. The returned value should
/// > only ever be used for statistics purposes.
pub fn total_outbound(&self) -> u64 {
self.inner.outbound.load(Ordering::Relaxed)
}
}
impl Default for BandwidthMeter {
fn default() -> Self {
Self {
inner: Arc::new(BandwidthMeterInner {
inbound: AtomicU64::new(0),
outbound: AtomicU64::new(0),
}),
}
}
}
/// Wraps around a single stream that implements [`AsyncRead`] + [`AsyncWrite`] and meters the
/// bandwidth through it
#[derive(Debug)]
#[pin_project::pin_project]
pub struct MeteredStream<S> {
/// The stream this instruments
#[pin]
inner: S,
/// The [`BandwidthMeter`] struct this uses to meter bandwidth
meter: BandwidthMeter,
}
impl<S> MeteredStream<S> {
/// Creates a new [`MeteredStream`] wrapping around the provided stream,
/// along with a new [`BandwidthMeter`]
pub fn new(inner: S) -> Self {
Self { inner, meter: BandwidthMeter::default() }
}
/// Creates a new [`MeteredStream`] wrapping around the provided stream,
/// attaching the provided [`BandwidthMeter`]
pub const fn new_with_meter(inner: S, meter: BandwidthMeter) -> Self {
Self { inner, meter }
}
/// Provides a reference to the [`BandwidthMeter`] attached to this [`MeteredStream`]
pub const fn get_bandwidth_meter(&self) -> &BandwidthMeter {
&self.meter
}
/// Returns the wrapped stream
pub const fn inner(&self) -> &S {
&self.inner
}
}
impl<Stream: AsyncRead> AsyncRead for MeteredStream<Stream> {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let this = self.project();
let num_bytes = {
let init_num_bytes = buf.filled().len();
ready!(this.inner.poll_read(cx, buf))?;
buf.filled().len() - init_num_bytes
};
this.meter
.inner
.inbound
.fetch_add(u64::try_from(num_bytes).unwrap_or(u64::MAX), Ordering::Relaxed);
Poll::Ready(Ok(()))
}
}
impl<Stream: AsyncWrite> AsyncWrite for MeteredStream<Stream> {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
let this = self.project();
let num_bytes = ready!(this.inner.poll_write(cx, buf))?;
this.meter
.inner
.outbound
.fetch_add(u64::try_from(num_bytes).unwrap_or(u64::MAX), Ordering::Relaxed);
Poll::Ready(Ok(num_bytes))
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
let this = self.project();
this.inner.poll_flush(cx)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
let this = self.project();
this.inner.poll_shutdown(cx)
}
}
impl HasRemoteAddr for MeteredStream<TcpStream> {
fn remote_addr(&self) -> Option<SocketAddr> {
self.inner.remote_addr()
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::{
io::{duplex, AsyncReadExt, AsyncWriteExt, DuplexStream},
net::TcpListener,
};
async fn duplex_stream_ping_pong(
client: &mut MeteredStream<DuplexStream>,
server: &mut MeteredStream<DuplexStream>,
) {
let mut buf = [0u8; 4];
client.write_all(b"ping").await.unwrap();
server.read_exact(&mut buf).await.unwrap();
server.write_all(b"pong").await.unwrap();
client.read_exact(&mut buf).await.unwrap();
}
fn assert_bandwidth_counts(
bandwidth_meter: &BandwidthMeter,
expected_inbound: u64,
expected_outbound: u64,
) {
let actual_inbound = bandwidth_meter.total_inbound();
assert_eq!(
actual_inbound, expected_inbound,
"Expected {expected_inbound} inbound bytes, but got {actual_inbound}",
);
let actual_outbound = bandwidth_meter.total_outbound();
assert_eq!(
actual_outbound, expected_outbound,
"Expected {expected_outbound} inbound bytes, but got {actual_outbound}",
);
}
#[tokio::test]
async fn test_count_read_write() {
// Taken in large part from https://docs.rs/tokio/latest/tokio/io/struct.DuplexStream.html#example
let (client, server) = duplex(64);
let mut metered_client = MeteredStream::new(client);
let mut metered_server = MeteredStream::new(server);
duplex_stream_ping_pong(&mut metered_client, &mut metered_server).await;
assert_bandwidth_counts(metered_client.get_bandwidth_meter(), 4, 4);
assert_bandwidth_counts(metered_server.get_bandwidth_meter(), 4, 4);
}
#[tokio::test]
async fn test_read_equals_write_tcp() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let server_addr = listener.local_addr().unwrap();
let client_stream = TcpStream::connect(server_addr).await.unwrap();
let mut metered_client_stream = MeteredStream::new(client_stream);
let client_meter = metered_client_stream.meter.clone();
let handle = tokio::spawn(async move {
let (server_stream, _) = listener.accept().await.unwrap();
let mut metered_server_stream = MeteredStream::new(server_stream);
let mut buf = [0u8; 4];
metered_server_stream.read_exact(&mut buf).await.unwrap();
assert_eq!(metered_server_stream.meter.total_inbound(), client_meter.total_outbound());
});
metered_client_stream.write_all(b"ping").await.unwrap();
handle.await.unwrap();
}
#[tokio::test]
async fn test_multiple_streams_one_meter() {
let (client_1, server_1) = duplex(64);
let (client_2, server_2) = duplex(64);
let shared_client_bandwidth_meter = BandwidthMeter::default();
let shared_server_bandwidth_meter = BandwidthMeter::default();
let mut metered_client_1 =
MeteredStream::new_with_meter(client_1, shared_client_bandwidth_meter.clone());
let mut metered_server_1 =
MeteredStream::new_with_meter(server_1, shared_server_bandwidth_meter.clone());
let mut metered_client_2 =
MeteredStream::new_with_meter(client_2, shared_client_bandwidth_meter.clone());
let mut metered_server_2 =
MeteredStream::new_with_meter(server_2, shared_server_bandwidth_meter.clone());
duplex_stream_ping_pong(&mut metered_client_1, &mut metered_server_1).await;
duplex_stream_ping_pong(&mut metered_client_2, &mut metered_server_2).await;
assert_bandwidth_counts(&shared_client_bandwidth_meter, 8, 8);
assert_bandwidth_counts(&shared_server_bandwidth_meter, 8, 8);
}
}

1
crates/net/common/src/lib.rs
View File

@ -9,7 +9,6 @@
#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
pub mod ban_list;
pub mod bandwidth_meter;
/// Traits related to tokio streams
pub mod stream;