From 256283966a6b932cb23e6a982a0fab5d74fbd293 Mon Sep 17 00:00:00 2001 From: AntiochP <30642645+antiochp@users.noreply.github.com> Date: Fri, 17 Nov 2017 20:13:49 -0500 Subject: [PATCH] Allowable write size as u64 instead of usize. Fix #303. --- p2p/src/rate_limit.rs | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/p2p/src/rate_limit.rs b/p2p/src/rate_limit.rs index 88cae361..d6998677 100644 --- a/p2p/src/rate_limit.rs +++ b/p2p/src/rate_limit.rs @@ -15,7 +15,7 @@ //! Provides wrappers for throttling readers and writers use std::time::Instant; -use std::io; +use std::{cmp, io}; use futures::*; use tokio_io::*; @@ -106,9 +106,9 @@ impl AsyncRead for ThrottledReader {} pub struct ThrottledWriter { writer: W, /// Max Bytes per second - max: u32, + max: u64, /// Stores a count of last request and last request time - allowed: usize, + allowed: u64, last_check: Instant, } @@ -116,11 +116,11 @@ pub struct ThrottledWriter { impl ThrottledWriter { /// Adds throttling to a writer. /// The resulting writer will write at most `max` amount of bytes per second - pub fn new(writer: W, max: u32) -> Self { + pub fn new(writer: W, max: u64) -> Self { ThrottledWriter { writer: writer, max: max, - allowed: max as usize, + allowed: max, last_check: Instant::now(), } } @@ -149,11 +149,11 @@ impl io::Write for ThrottledWriter { // Check passed Time let time_passed = self.last_check.elapsed(); self.last_check = Instant::now(); - self.allowed += time_passed.as_secs() as usize * self.max as usize; + self.allowed += time_passed.as_secs() * self.max; // Throttle - if self.allowed > self.max as usize { - self.allowed = self.max as usize; + if self.allowed > self.max { + self.allowed = cmp::min(self.max, u32::max_value() as u64); } // Check if Allowed @@ -165,8 +165,8 @@ impl io::Write for ThrottledWriter { } // Write max allowed - let buf = if buf.len() > self.allowed { - &buf[0..self.allowed] + let buf = if buf.len() > self.allowed as usize { + &buf[0..self.allowed as usize] } else { buf }; @@ -174,7 +174,7 @@ impl io::Write for ThrottledWriter { // Decrement Allowed amount written if let Ok(n) = res { - self.allowed -= n; + self.allowed -= n as u64; } res }