Proxmox-Port/proxmox
1
0
Fork 0
forked from proxmox-mirrors/proxmox
Code Pull Requests Activity

http: rustfmt

Browse Source 
Signed-off-by: Thomas Lamprecht <t.lamprecht@proxmox.com>
This commit is contained in:
Thomas Lamprecht 2022-04-10 12:41:21 +02:00
parent 05cad8926b
commit 0eeb0dd17c
3 changed files with 64 additions and 66 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
proxmox-http/src/client/connector.rs
View File

@ -148,7 +148,6 @@ impl hyper::service::Service<Uri> for HttpsConnector {
let read_limiter = self.read_limiter.clone();
let write_limiter = self.write_limiter.clone();
if let Some(ref proxy) = self.proxy {
let use_connect = is_https || proxy.force_connect;
@ -177,11 +176,8 @@ impl hyper::service::Service<Uri> for HttpsConnector {
let _ = set_tcp_keepalive(tcp_stream.as_raw_fd(), keepalive);
let mut tcp_stream = RateLimitedStream::with_limiter(
tcp_stream,
read_limiter,
write_limiter,
);
let mut tcp_stream =
RateLimitedStream::with_limiter(tcp_stream, read_limiter, write_limiter);
let mut connect_request = format!("CONNECT {0}:{1} HTTP/1.1\r\n", host, port);
if let Some(authorization) = authorization {
@ -210,11 +206,8 @@ impl hyper::service::Service<Uri> for HttpsConnector {
let _ = set_tcp_keepalive(tcp_stream.as_raw_fd(), keepalive);
let tcp_stream = RateLimitedStream::with_limiter(
tcp_stream,
read_limiter,
write_limiter,
);
let tcp_stream =
RateLimitedStream::with_limiter(tcp_stream, read_limiter, write_limiter);
Ok(MaybeTlsStream::Proxied(tcp_stream))
}
@ -230,11 +223,8 @@ impl hyper::service::Service<Uri> for HttpsConnector {
let _ = set_tcp_keepalive(tcp_stream.as_raw_fd(), keepalive);
let tcp_stream = RateLimitedStream::with_limiter(
tcp_stream,
read_limiter,
write_limiter,
);
let tcp_stream =
RateLimitedStream::with_limiter(tcp_stream, read_limiter, write_limiter);
if is_https {
Self::secure_stream(tcp_stream, &ssl_connector, &host).await

62
proxmox-http/src/client/rate_limited_stream.rs
View File

@ -1,17 +1,17 @@
use std::pin::Pin;
use std::io::IoSlice;
use std::marker::Unpin;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use std::io::IoSlice;
use futures::Future;
use tokio::io::{ReadBuf, AsyncRead, AsyncWrite};
use hyper::client::connect::{Connected, Connection};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::time::Sleep;
use hyper::client::connect::{Connection, Connected};
use std::task::{Context, Poll};
use super::{ShareableRateLimit, RateLimiter};
use super::{RateLimiter, ShareableRateLimit};
type SharedRateLimit = Arc<dyn ShareableRateLimit>;
@ -21,7 +21,8 @@ pub struct RateLimitedStream<S> {
read_delay: Option<Pin<Box<Sleep>>>,
write_limiter: Option<SharedRateLimit>,
write_delay: Option<Pin<Box<Sleep>>>,
update_limiter_cb: Option<Box<dyn Fn() -> (Option<SharedRateLimit>, Option<SharedRateLimit>) + Send>>,
update_limiter_cb:
Option<Box<dyn Fn() -> (Option<SharedRateLimit>, Option<SharedRateLimit>) + Send>>,
last_limiter_update: Instant,
stream: S,
}
@ -32,8 +33,7 @@ impl RateLimitedStream<tokio::net::TcpStream> {
}
}
impl <S> RateLimitedStream<S> {
impl<S> RateLimitedStream<S> {
/// Creates a new instance with reads and writes limited to the same `rate`.
pub fn new(stream: S, rate: u64, bucket_size: u64) -> Self {
let now = Instant::now();
@ -50,7 +50,7 @@ impl <S> RateLimitedStream<S> {
read_limiter: Option<SharedRateLimit>,
write_limiter: Option<SharedRateLimit>,
) -> Self {
Self {
Self {
read_limiter,
read_delay: None,
write_limiter,
@ -67,7 +67,9 @@ impl <S> RateLimitedStream<S> {
///
/// Note: This function is called within an async context, so it
/// should be fast and must not block.
pub fn with_limiter_update_cb<F: Fn() -> (Option<SharedRateLimit>, Option<SharedRateLimit>) + Send + 'static>(
pub fn with_limiter_update_cb<
F: Fn() -> (Option<SharedRateLimit>, Option<SharedRateLimit>) + Send + 'static,
>(
stream: S,
update_limiter_cb: F,
) -> Self {
@ -95,11 +97,7 @@ impl <S> RateLimitedStream<S> {
}
}
fn register_traffic(
limiter: &(dyn ShareableRateLimit),
count: usize,
) -> Option<Pin<Box<Sleep>>>{
fn register_traffic(limiter: &(dyn ShareableRateLimit), count: usize) -> Option<Pin<Box<Sleep>>> {
const MIN_DELAY: Duration = Duration::from_millis(10);
let now = Instant::now();
@ -114,25 +112,24 @@ fn register_traffic(
fn delay_is_ready(delay: &mut Option<Pin<Box<Sleep>>>, ctx: &mut Context<'_>) -> bool {
match delay {
Some(ref mut future) => {
future.as_mut().poll(ctx).is_ready()
}
Some(ref mut future) => future.as_mut().poll(ctx).is_ready(),
None => true,
}
}
impl <S: AsyncWrite + Unpin> AsyncWrite for RateLimitedStream<S> {
impl<S: AsyncWrite + Unpin> AsyncWrite for RateLimitedStream<S> {
fn poll_write(
self: Pin<&mut Self>,
ctx: &mut Context<'_>,
buf: &[u8]
buf: &[u8],
) -> Poll<Result<usize, std::io::Error>> {
let this = self.get_mut();
let is_ready = delay_is_ready(&mut this.write_delay, ctx);
if !is_ready { return Poll::Pending; }
if !is_ready {
return Poll::Pending;
}
this.write_delay = None;
@ -156,13 +153,15 @@ impl <S: AsyncWrite + Unpin> AsyncWrite for RateLimitedStream<S> {
fn poll_write_vectored(
self: Pin<&mut Self>,
ctx: &mut Context<'_>,
bufs: &[IoSlice<'_>]
bufs: &[IoSlice<'_>],
) -> Poll<Result<usize, std::io::Error>> {
let this = self.get_mut();
let is_ready = delay_is_ready(&mut this.write_delay, ctx);
if !is_ready { return Poll::Pending; }
if !is_ready {
return Poll::Pending;
}
this.write_delay = None;
@ -179,25 +178,21 @@ impl <S: AsyncWrite + Unpin> AsyncWrite for RateLimitedStream<S> {
result
}
fn poll_flush(
self: Pin<&mut Self>,
ctx: &mut Context<'_>
) -> Poll<Result<(), std::io::Error>> {
fn poll_flush(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> {
let this = self.get_mut();
Pin::new(&mut this.stream).poll_flush(ctx)
}
fn poll_shutdown(
self: Pin<&mut Self>,
ctx: &mut Context<'_>
ctx: &mut Context<'_>,
) -> Poll<Result<(), std::io::Error>> {
let this = self.get_mut();
Pin::new(&mut this.stream).poll_shutdown(ctx)
}
}
impl <S: AsyncRead + Unpin> AsyncRead for RateLimitedStream<S> {
impl<S: AsyncRead + Unpin> AsyncRead for RateLimitedStream<S> {
fn poll_read(
self: Pin<&mut Self>,
ctx: &mut Context<'_>,
@ -207,7 +202,9 @@ impl <S: AsyncRead + Unpin> AsyncRead for RateLimitedStream<S> {
let is_ready = delay_is_ready(&mut this.read_delay, ctx);
if !is_ready { return Poll::Pending; }
if !is_ready {
return Poll::Pending;
}
this.read_delay = None;
@ -225,7 +222,6 @@ impl <S: AsyncRead + Unpin> AsyncRead for RateLimitedStream<S> {
result
}
}
// we need this for the hyper http client

46
proxmox-http/src/client/rate_limiter.rs
View File

@ -1,5 +1,5 @@
use std::time::{Duration, Instant};
use std::convert::TryInto;
use std::time::{Duration, Instant};
use anyhow::{bail, Error};
@ -36,7 +36,6 @@ struct TbfState {
}
impl TbfState {
const NO_DELAY: Duration = Duration::from_millis(0);
fn refill_bucket(&mut self, rate: u64, current_time: Instant) {
@ -45,12 +44,15 @@ impl TbfState {
None => return,
};
if time_diff == 0 { return; }
if time_diff == 0 {
return;
}
self.last_update = current_time;
let allowed_traffic = ((time_diff.saturating_mul(rate as u128)) / 1_000_000_000)
.try_into().unwrap_or(u64::MAX);
.try_into()
.unwrap_or(u64::MAX);
self.consumed_tokens = self.consumed_tokens.saturating_sub(allowed_traffic);
}
@ -70,7 +72,9 @@ impl TbfState {
if self.consumed_tokens <= bucket_size {
return Self::NO_DELAY;
}
Duration::from_nanos((self.consumed_tokens - bucket_size).saturating_mul(1_000_000_000)/rate)
Duration::from_nanos(
(self.consumed_tokens - bucket_size).saturating_mul(1_000_000_000) / rate,
)
}
}
@ -80,13 +84,12 @@ impl TbfState {
/// change/modify the layout (do not add fields)
#[repr(C)]
pub struct RateLimiter {
rate: u64, // tokens/second
rate: u64, // tokens/second
bucket_size: u64, // TBF bucket size
state: TbfState,
}
impl RateLimiter {
/// Creates a new instance, using [Instant::now] as start time.
pub fn new(rate: u64, bucket_size: u64) -> Self {
let start_time = Instant::now();
@ -109,7 +112,6 @@ impl RateLimiter {
}
impl RateLimit for RateLimiter {
fn update_rate(&mut self, rate: u64, bucket_size: u64) {
self.rate = rate;
@ -125,12 +127,12 @@ impl RateLimit for RateLimiter {
}
fn register_traffic(&mut self, current_time: Instant, data_len: u64) -> Duration {
self.state.register_traffic(self.rate, self.bucket_size, current_time, data_len)
self.state
.register_traffic(self.rate, self.bucket_size, current_time, data_len)
}
}
impl <R: RateLimit + Send> ShareableRateLimit for std::sync::Mutex<R> {
impl<R: RateLimit + Send> ShareableRateLimit for std::sync::Mutex<R> {
fn update_rate(&self, rate: u64, bucket_size: u64) {
self.lock().unwrap().update_rate(rate, bucket_size);
}
@ -140,22 +142,22 @@ impl <R: RateLimit + Send> ShareableRateLimit for std::sync::Mutex<R> {
}
fn register_traffic(&self, current_time: Instant, data_len: u64) -> Duration {
self.lock().unwrap().register_traffic(current_time, data_len)
self.lock()
.unwrap()
.register_traffic(current_time, data_len)
}
}
/// Array of rate limiters.
///
/// A group of rate limiters with same configuration.
pub struct RateLimiterVec {
rate: u64, // tokens/second
rate: u64, // tokens/second
bucket_size: u64, // TBF bucket size
state: Vec<TbfState>,
}
impl RateLimiterVec {
/// Creates a new instance, using [Instant::now] as start time.
pub fn new(group_size: usize, rate: u64, bucket_size: u64) -> Self {
let start_time = Instant::now();
@ -163,7 +165,12 @@ impl RateLimiterVec {
}
/// Creates a new instance with specified `rate`, `bucket_size` and `start_time`.
pub fn with_start_time(group_size: usize, rate: u64, bucket_size: u64, start_time: Instant) -> Self {
pub fn with_start_time(
group_size: usize,
rate: u64,
bucket_size: u64,
start_time: Instant,
) -> Self {
let state = TbfState {
traffic: 0,
last_update: start_time,
@ -191,7 +198,12 @@ impl RateLimiterVec {
}
/// Register traffic at the specified index
pub fn register_traffic(&mut self, index: usize, current_time: Instant, data_len: u64) -> Result<Duration, Error> {
pub fn register_traffic(
&mut self,
index: usize,
current_time: Instant,
data_len: u64,
) -> Result<Duration, Error> {
if index >= self.state.len() {
bail!("RateLimiterVec::register_traffic - index out of range");
}