gringron_util/
rate_counter.rs

1// Copyright 2021 The GrinGron 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
15use std::convert::TryInto;
16/// Utility to track the rate of data transfers
17use std::time::SystemTime;
18
19struct Entry {
20	bytes: u64,
21	timestamp: u64,
22}
23
24impl Entry {
25	fn new(bytes: u64) -> Entry {
26		Entry {
27			bytes,
28			timestamp: millis_since_epoch(),
29		}
30	}
31
32	// Create new "quiet" entry with zero timestamp.
33	// This will count toward total bytes but will not affect the "msg rate".
34	fn new_quiet(bytes: u64) -> Entry {
35		Entry {
36			bytes,
37			timestamp: 0,
38		}
39	}
40
41	// We want to filter out "quiet" entries when calculating the "msg rate".
42	fn is_quiet(&self) -> bool {
43		self.timestamp == 0
44	}
45}
46
47/// A rate counter tracks the number of transfers, the amount of data
48/// exchanged and the rate of transfer (via a few timers) over the last
49/// minute. The counter does not try to be accurate and update times
50/// proactively, instead it only does so lazily. As a result, produced
51/// rates are worst-case estimates.
52pub struct RateCounter {
53	last_min_entries: Vec<Entry>,
54}
55
56impl RateCounter {
57	/// Instantiate a new rate counter
58	pub fn new() -> RateCounter {
59		RateCounter {
60			last_min_entries: vec![],
61		}
62	}
63
64	/// Increments number of bytes transferred, updating counts and rates.
65	pub fn inc(&mut self, bytes: u64) {
66		self.last_min_entries.push(Entry::new(bytes));
67		self.truncate();
68	}
69
70	/// Increments number of bytes without updating the count or rate.
71	/// We filter out 0 last_min_times when calculating rate.
72	/// Used during txhashset.zip download to track bytes downloaded
73	/// without treating a peer as abusive (too high a rate of download).
74	pub fn inc_quiet(&mut self, bytes: u64) {
75		self.last_min_entries.push(Entry::new_quiet(bytes));
76		self.truncate();
77	}
78
79	fn truncate(&mut self) {
80		let now_millis = millis_since_epoch();
81		while !self.last_min_entries.is_empty()
82			&& self.last_min_entries[0].timestamp + 60000 < now_millis
83		{
84			self.last_min_entries.remove(0);
85		}
86	}
87
88	/// Number of bytes counted in the last minute.
89	/// Includes "quiet" byte increments.
90	pub fn bytes_per_min(&self) -> u64 {
91		self.last_min_entries.iter().map(|x| x.bytes).sum()
92	}
93
94	/// Count of increases in the last minute.
95	/// Excludes "quiet" byte increments.
96	pub fn count_per_min(&self) -> u64 {
97		self.last_min_entries
98			.iter()
99			.filter(|x| !x.is_quiet())
100			.count() as u64
101	}
102
103	/// Elapsed time in ms since the last entry.
104	/// We use this to rate limit when sending.
105	pub fn elapsed_since_last_msg(&self) -> Option<u64> {
106		self.last_min_entries
107			.last()
108			.map(|x| millis_since_epoch().saturating_sub(x.timestamp))
109	}
110}
111
112// turns out getting the millisecs since epoch in Rust isn't as easy as it
113// could be
114fn millis_since_epoch() -> u64 {
115	SystemTime::now()
116		.duration_since(SystemTime::UNIX_EPOCH)
117		.map(|since_epoch| since_epoch.as_millis().try_into().unwrap_or(0))
118		.unwrap_or(0)
119}
gringron_util/rate_counter.rs

gringron_util/
rate_counter.rs
