github.com/vipernet-xyz/tm@v0.34.24/libs/flowrate/flowrate.go (about) 1 // 2 // Written by Maxim Khitrov (November 2012) 3 // 4 5 // Package flowrate provides the tools for monitoring and limiting the flow rate 6 // of an arbitrary data stream. 7 package flowrate 8 9 import ( 10 "math" 11 "time" 12 13 tmsync "github.com/vipernet-xyz/tm/libs/sync" 14 ) 15 16 // Monitor monitors and limits the transfer rate of a data stream. 17 type Monitor struct { 18 mu tmsync.Mutex // Mutex guarding access to all internal fields 19 active bool // Flag indicating an active transfer 20 start time.Duration // Transfer start time (clock() value) 21 bytes int64 // Total number of bytes transferred 22 samples int64 // Total number of samples taken 23 24 rSample float64 // Most recent transfer rate sample (bytes per second) 25 rEMA float64 // Exponential moving average of rSample 26 rPeak float64 // Peak transfer rate (max of all rSamples) 27 rWindow float64 // rEMA window (seconds) 28 29 sBytes int64 // Number of bytes transferred since sLast 30 sLast time.Duration // Most recent sample time (stop time when inactive) 31 sRate time.Duration // Sampling rate 32 33 tBytes int64 // Number of bytes expected in the current transfer 34 tLast time.Duration // Time of the most recent transfer of at least 1 byte 35 } 36 37 // New creates a new flow control monitor. Instantaneous transfer rate is 38 // measured and updated for each sampleRate interval. windowSize determines the 39 // weight of each sample in the exponential moving average (EMA) calculation. 40 // The exact formulas are: 41 // 42 // sampleTime = currentTime - prevSampleTime 43 // sampleRate = byteCount / sampleTime 44 // weight = 1 - exp(-sampleTime/windowSize) 45 // newRate = weight*sampleRate + (1-weight)*oldRate 46 // 47 // The default values for sampleRate and windowSize (if <= 0) are 100ms and 1s, 48 // respectively. 49 func New(sampleRate, windowSize time.Duration) *Monitor { 50 if sampleRate = clockRound(sampleRate); sampleRate <= 0 { 51 sampleRate = 5 * clockRate 52 } 53 if windowSize <= 0 { 54 windowSize = 1 * time.Second 55 } 56 now := clock() 57 return &Monitor{ 58 active: true, 59 start: now, 60 rWindow: windowSize.Seconds(), 61 sLast: now, 62 sRate: sampleRate, 63 tLast: now, 64 } 65 } 66 67 // Update records the transfer of n bytes and returns n. It should be called 68 // after each Read/Write operation, even if n is 0. 69 func (m *Monitor) Update(n int) int { 70 m.mu.Lock() 71 m.update(n) 72 m.mu.Unlock() 73 return n 74 } 75 76 // Hack to set the current rEMA. 77 func (m *Monitor) SetREMA(rEMA float64) { 78 m.mu.Lock() 79 m.rEMA = rEMA 80 m.samples++ 81 m.mu.Unlock() 82 } 83 84 // IO is a convenience method intended to wrap io.Reader and io.Writer method 85 // execution. It calls m.Update(n) and then returns (n, err) unmodified. 86 func (m *Monitor) IO(n int, err error) (int, error) { 87 return m.Update(n), err 88 } 89 90 // Done marks the transfer as finished and prevents any further updates or 91 // limiting. Instantaneous and current transfer rates drop to 0. Update, IO, and 92 // Limit methods become NOOPs. It returns the total number of bytes transferred. 93 func (m *Monitor) Done() int64 { 94 m.mu.Lock() 95 if now := m.update(0); m.sBytes > 0 { 96 m.reset(now) 97 } 98 m.active = false 99 m.tLast = 0 100 n := m.bytes 101 m.mu.Unlock() 102 return n 103 } 104 105 // timeRemLimit is the maximum Status.TimeRem value. 106 const timeRemLimit = 999*time.Hour + 59*time.Minute + 59*time.Second 107 108 // Status represents the current Monitor status. All transfer rates are in bytes 109 // per second rounded to the nearest byte. 110 type Status struct { 111 Start time.Time // Transfer start time 112 Bytes int64 // Total number of bytes transferred 113 Samples int64 // Total number of samples taken 114 InstRate int64 // Instantaneous transfer rate 115 CurRate int64 // Current transfer rate (EMA of InstRate) 116 AvgRate int64 // Average transfer rate (Bytes / Duration) 117 PeakRate int64 // Maximum instantaneous transfer rate 118 BytesRem int64 // Number of bytes remaining in the transfer 119 Duration time.Duration // Time period covered by the statistics 120 Idle time.Duration // Time since the last transfer of at least 1 byte 121 TimeRem time.Duration // Estimated time to completion 122 Progress Percent // Overall transfer progress 123 Active bool // Flag indicating an active transfer 124 } 125 126 // Status returns current transfer status information. The returned value 127 // becomes static after a call to Done. 128 func (m *Monitor) Status() Status { 129 m.mu.Lock() 130 now := m.update(0) 131 s := Status{ 132 Active: m.active, 133 Start: clockToTime(m.start), 134 Duration: m.sLast - m.start, 135 Idle: now - m.tLast, 136 Bytes: m.bytes, 137 Samples: m.samples, 138 PeakRate: round(m.rPeak), 139 BytesRem: m.tBytes - m.bytes, 140 Progress: percentOf(float64(m.bytes), float64(m.tBytes)), 141 } 142 if s.BytesRem < 0 { 143 s.BytesRem = 0 144 } 145 if s.Duration > 0 { 146 rAvg := float64(s.Bytes) / s.Duration.Seconds() 147 s.AvgRate = round(rAvg) 148 if s.Active { 149 s.InstRate = round(m.rSample) 150 s.CurRate = round(m.rEMA) 151 if s.BytesRem > 0 { 152 if tRate := 0.8*m.rEMA + 0.2*rAvg; tRate > 0 { 153 ns := float64(s.BytesRem) / tRate * 1e9 154 if ns > float64(timeRemLimit) { 155 ns = float64(timeRemLimit) 156 } 157 s.TimeRem = clockRound(time.Duration(ns)) 158 } 159 } 160 } 161 } 162 m.mu.Unlock() 163 return s 164 } 165 166 // Limit restricts the instantaneous (per-sample) data flow to rate bytes per 167 // second. It returns the maximum number of bytes (0 <= n <= want) that may be 168 // transferred immediately without exceeding the limit. If block == true, the 169 // call blocks until n > 0. want is returned unmodified if want < 1, rate < 1, 170 // or the transfer is inactive (after a call to Done). 171 // 172 // At least one byte is always allowed to be transferred in any given sampling 173 // period. Thus, if the sampling rate is 100ms, the lowest achievable flow rate 174 // is 10 bytes per second. 175 // 176 // For usage examples, see the implementation of Reader and Writer in io.go. 177 func (m *Monitor) Limit(want int, rate int64, block bool) (n int) { 178 if want < 1 || rate < 1 { 179 return want 180 } 181 m.mu.Lock() 182 183 // Determine the maximum number of bytes that can be sent in one sample 184 limit := round(float64(rate) * m.sRate.Seconds()) 185 if limit <= 0 { 186 limit = 1 187 } 188 189 // If block == true, wait until m.sBytes < limit 190 if now := m.update(0); block { 191 for m.sBytes >= limit && m.active { 192 now = m.waitNextSample(now) 193 } 194 } 195 196 // Make limit <= want (unlimited if the transfer is no longer active) 197 if limit -= m.sBytes; limit > int64(want) || !m.active { 198 limit = int64(want) 199 } 200 m.mu.Unlock() 201 202 if limit < 0 { 203 limit = 0 204 } 205 return int(limit) 206 } 207 208 // SetTransferSize specifies the total size of the data transfer, which allows 209 // the Monitor to calculate the overall progress and time to completion. 210 func (m *Monitor) SetTransferSize(bytes int64) { 211 if bytes < 0 { 212 bytes = 0 213 } 214 m.mu.Lock() 215 m.tBytes = bytes 216 m.mu.Unlock() 217 } 218 219 // update accumulates the transferred byte count for the current sample until 220 // clock() - m.sLast >= m.sRate. The monitor status is updated once the current 221 // sample is done. 222 func (m *Monitor) update(n int) (now time.Duration) { 223 if !m.active { 224 return 225 } 226 if now = clock(); n > 0 { 227 m.tLast = now 228 } 229 m.sBytes += int64(n) 230 if sTime := now - m.sLast; sTime >= m.sRate { 231 t := sTime.Seconds() 232 if m.rSample = float64(m.sBytes) / t; m.rSample > m.rPeak { 233 m.rPeak = m.rSample 234 } 235 236 // Exponential moving average using a method similar to *nix load 237 // average calculation. Longer sampling periods carry greater weight. 238 if m.samples > 0 { 239 w := math.Exp(-t / m.rWindow) 240 m.rEMA = m.rSample + w*(m.rEMA-m.rSample) 241 } else { 242 m.rEMA = m.rSample 243 } 244 m.reset(now) 245 } 246 return 247 } 248 249 // reset clears the current sample state in preparation for the next sample. 250 func (m *Monitor) reset(sampleTime time.Duration) { 251 m.bytes += m.sBytes 252 m.samples++ 253 m.sBytes = 0 254 m.sLast = sampleTime 255 } 256 257 // waitNextSample sleeps for the remainder of the current sample. The lock is 258 // released and reacquired during the actual sleep period, so it's possible for 259 // the transfer to be inactive when this method returns. 260 func (m *Monitor) waitNextSample(now time.Duration) time.Duration { 261 const minWait = 5 * time.Millisecond 262 current := m.sLast 263 264 // sleep until the last sample time changes (ideally, just one iteration) 265 for m.sLast == current && m.active { 266 d := current + m.sRate - now 267 m.mu.Unlock() 268 if d < minWait { 269 d = minWait 270 } 271 time.Sleep(d) 272 m.mu.Lock() 273 now = m.update(0) 274 } 275 return now 276 }