From mboxrd@z Thu Jan 1 00:00:00 1970 Return-Path: Received: from firstgate.proxmox.com (firstgate.proxmox.com [212.224.123.68]) (using TLSv1.3 with cipher TLS_AES_256_GCM_SHA384 (256/256 bits) key-exchange X25519 server-signature RSA-PSS (2048 bits)) (No client certificate requested) by lists.proxmox.com (Postfix) with ESMTPS id 686C57D501 for ; Tue, 9 Nov 2021 07:53:05 +0100 (CET) Received: from firstgate.proxmox.com (localhost [127.0.0.1]) by firstgate.proxmox.com (Proxmox) with ESMTP id A5C162DFF9 for ; Tue, 9 Nov 2021 07:53:04 +0100 (CET) Received: from proxmox-new.maurer-it.com (proxmox-new.maurer-it.com [94.136.29.106]) (using TLSv1.3 with cipher TLS_AES_256_GCM_SHA384 (256/256 bits) key-exchange X25519 server-signature RSA-PSS (2048 bits) server-digest SHA256) (No client certificate requested) by firstgate.proxmox.com (Proxmox) with ESMTPS id 5E7E12DF59 for ; Tue, 9 Nov 2021 07:53:01 +0100 (CET) Received: from proxmox-new.maurer-it.com (localhost.localdomain [127.0.0.1]) by proxmox-new.maurer-it.com (Proxmox) with ESMTP id 7E7D542CF2; Tue, 9 Nov 2021 07:52:55 +0100 (CET) From: Dietmar Maurer To: pbs-devel@lists.proxmox.com Date: Tue, 9 Nov 2021 07:52:40 +0100 Message-Id: <20211109065253.980304-4-dietmar@proxmox.com> X-Mailer: git-send-email 2.30.2 In-Reply-To: <20211109065253.980304-1-dietmar@proxmox.com> References: <20211109065253.980304-1-dietmar@proxmox.com> MIME-Version: 1.0 Content-Transfer-Encoding: 8bit X-SPAM-LEVEL: Spam detection results: 0 AWL 0.514 Adjusted score from AWL reputation of From: address BAYES_00 -1.9 Bayes spam probability is 0 to 1% KAM_DMARC_STATUS 0.01 Test Rule for DKIM or SPF Failure with Strict Alignment SPF_HELO_NONE 0.001 SPF: HELO does not publish an SPF Record SPF_PASS -0.001 SPF: sender matches SPF record Subject: [pbs-devel] [PATCH proxmox 2/7] RateLimitedStream: implement poll_write_vectored X-BeenThere: pbs-devel@lists.proxmox.com X-Mailman-Version: 2.1.29 Precedence: list List-Id: Proxmox Backup Server development discussion List-Unsubscribe: , List-Archive: List-Post: List-Help: List-Subscribe: , X-List-Received-Date: Tue, 09 Nov 2021 06:53:05 -0000 Signed-off-by: Dietmar Maurer --- .../src/client/rate_limited_stream.rs | 92 +++++++++++++------ 1 file changed, 62 insertions(+), 30 deletions(-) diff --git a/proxmox-http/src/client/rate_limited_stream.rs b/proxmox-http/src/client/rate_limited_stream.rs index a11b59e..0cc0ebb 100644 --- a/proxmox-http/src/client/rate_limited_stream.rs +++ b/proxmox-http/src/client/rate_limited_stream.rs @@ -2,6 +2,7 @@ use std::pin::Pin; use std::marker::Unpin; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; +use std::io::IoSlice; use futures::Future; use tokio::io::{ReadBuf, AsyncRead, AsyncWrite}; @@ -22,8 +23,6 @@ pub struct RateLimitedStream { impl RateLimitedStream { - const MIN_DELAY: Duration = Duration::from_millis(20); - /// 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(); @@ -48,6 +47,33 @@ impl RateLimitedStream { } } +fn register_traffic( + limiter: &Mutex, + count: usize, +) -> Option>>{ + + const MIN_DELAY: Duration = Duration::from_millis(10); + + let now = Instant::now(); + let delay = limiter.lock().unwrap() + .register_traffic(now, count as u64); + if delay >= MIN_DELAY { + let sleep = tokio::time::sleep(delay); + Some(Box::pin(sleep)) + } else { + None + } +} + +fn delay_is_ready(delay: &mut Option>>, ctx: &mut Context<'_>) -> bool { + match delay { + Some(ref mut future) => { + future.as_mut().poll(ctx).is_ready() + } + None => true, + } +} + impl AsyncWrite for RateLimitedStream { fn poll_write( @@ -57,12 +83,7 @@ impl AsyncWrite for RateLimitedStream { ) -> Poll> { let this = self.get_mut(); - let is_ready = match this.write_delay { - Some(ref mut future) => { - future.as_mut().poll(ctx).is_ready() - } - None => true, - }; + let is_ready = delay_is_ready(&mut this.write_delay, ctx); if !is_ready { return Poll::Pending; } @@ -70,15 +91,37 @@ impl AsyncWrite for RateLimitedStream { let result = Pin::new(&mut this.stream).poll_write(ctx, buf); - if let Some(ref write_limiter) = this.write_limiter { - if let Poll::Ready(Ok(count)) = &result { - let now = Instant::now(); - let delay = write_limiter.lock().unwrap() - .register_traffic(now, *count as u64); - if delay >= Self::MIN_DELAY { - let sleep = tokio::time::sleep(delay); - this.write_delay = Some(Box::pin(sleep)); - } + if let Some(ref limiter) = this.write_limiter { + if let Poll::Ready(Ok(count)) = result { + this.write_delay = register_traffic(limiter, count); + } + } + + result + } + + fn is_write_vectored(&self) -> bool { + self.stream.is_write_vectored() + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + ctx: &mut Context<'_>, + bufs: &[IoSlice<'_>] + ) -> Poll> { + let this = self.get_mut(); + + let is_ready = delay_is_ready(&mut this.write_delay, ctx); + + if !is_ready { return Poll::Pending; } + + this.write_delay = None; + + let result = Pin::new(&mut this.stream).poll_write_vectored(ctx, bufs); + + if let Some(ref limiter) = this.write_limiter { + if let Poll::Ready(Ok(count)) = result { + this.write_delay = register_traffic(limiter, count); } } @@ -111,12 +154,7 @@ impl AsyncRead for RateLimitedStream { ) -> Poll> { let this = self.get_mut(); - let is_ready = match this.read_delay { - Some(ref mut future) => { - future.as_mut().poll(ctx).is_ready() - } - None => true, - }; + let is_ready = delay_is_ready(&mut this.read_delay, ctx); if !is_ready { return Poll::Pending; } @@ -128,13 +166,7 @@ impl AsyncRead for RateLimitedStream { if let Some(ref read_limiter) = this.read_limiter { if let Poll::Ready(Ok(())) = &result { let count = buf.filled().len() - filled_len; - let now = Instant::now(); - let delay = read_limiter.lock().unwrap() - .register_traffic(now, count as u64); - if delay >= Self::MIN_DELAY { - let sleep = tokio::time::sleep(delay); - this.read_delay = Some(Box::pin(sleep)); - } + this.read_delay = register_traffic(read_limiter, count); } } -- 2.30.2