stack_test_epic_util/
rate_counter.rs

1// Copyright 2018 The Grin Developers
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15/// Utility to track the rate of data transfers
16use std::time::{Duration, SystemTime};
17
18struct Entry {
19	bytes: u64,
20	timestamp: u64,
21}
22
23impl Entry {
24	fn new(bytes: u64) -> Entry {
25		Entry {
26			bytes,
27			timestamp: millis_since_epoch(),
28		}
29	}
30
31	// Create new "quiet" entry with zero timestamp.
32	// This will count toward total bytes but will not affect the "msg rate".
33	fn new_quiet(bytes: u64) -> Entry {
34		Entry {
35			bytes,
36			timestamp: 0,
37		}
38	}
39
40	// We want to filter out "quiet" entries when calculating the "msg rate".
41	fn is_quiet(&self) -> bool {
42		self.timestamp == 0
43	}
44}
45
46/// A rate counter tracks the number of transfers, the amount of data
47/// exchanged and the rate of transfer (via a few timers) over the last
48/// minute. The counter does not try to be accurate and update times
49/// proactively, instead it only does so lazily. As a result, produced
50/// rates are worst-case estimates.
51pub struct RateCounter {
52	last_min_entries: Vec<Entry>,
53}
54
55impl RateCounter {
56	/// Instantiate a new rate counter
57	pub fn new() -> RateCounter {
58		RateCounter {
59			last_min_entries: vec![],
60		}
61	}
62
63	/// Increments number of bytes transferred, updating counts and rates.
64	pub fn inc(&mut self, bytes: u64) {
65		self.last_min_entries.push(Entry::new(bytes));
66		self.truncate();
67	}
68
69	/// Increments number of bytes without updating the count or rate.
70	/// We filter out 0 last_min_times when calculating rate.
71	/// Used during txhashset.zip download to track bytes downloaded
72	/// without treating a peer as abusive (too high a rate of download).
73	pub fn inc_quiet(&mut self, bytes: u64) {
74		self.last_min_entries.push(Entry::new_quiet(bytes));
75		self.truncate();
76	}
77
78	fn truncate(&mut self) {
79		let now_millis = millis_since_epoch();
80		while self.last_min_entries.len() > 0
81			&& self.last_min_entries[0].timestamp + 60000 < now_millis
82		{
83			self.last_min_entries.remove(0);
84		}
85	}
86
87	/// Number of bytes counted in the last minute.
88	/// Includes "quiet" byte increments.
89	pub fn bytes_per_min(&self) -> u64 {
90		self.last_min_entries.iter().map(|x| x.bytes).sum()
91	}
92
93	/// Count of increases in the last minute.
94	/// Excludes "quiet" byte increments.
95	pub fn count_per_min(&self) -> u64 {
96		self.last_min_entries
97			.iter()
98			.filter(|x| !x.is_quiet())
99			.count() as u64
100	}
101}
102
103// turns out getting the millisecs since epoch in Rust isn't as easy as it
104// could be
105fn millis_since_epoch() -> u64 {
106	let since_epoch = SystemTime::now()
107		.duration_since(SystemTime::UNIX_EPOCH)
108		.unwrap_or(Duration::new(0, 0));
109	since_epoch.as_secs() * 1000 + since_epoch.subsec_millis() as u64
110}