github.com/badrootd/celestia-core@v0.0.0-20240305091328-aa4207a4b25d/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 cmtsync "github.com/badrootd/celestia-core/libs/sync" 14 ) 15 16 // Monitor monitors and limits the transfer rate of a data stream. 17 type Monitor struct { 18 mu cmtsync.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 // GetSampleRate returns the current sampling rate. 68 func (m *Monitor) GetSampleRate() time.Duration { 69 m.mu.Lock() 70 defer m.mu.Unlock() 71 return m.sRate 72 } 73 74 // Update records the transfer of n bytes and returns n. It should be called 75 // after each Read/Write operation, even if n is 0. 76 func (m *Monitor) Update(n int) int { 77 m.mu.Lock() 78 m.update(n) 79 m.mu.Unlock() 80 return n 81 } 82 83 // Hack to set the current rEMA. 84 func (m *Monitor) SetREMA(rEMA float64) { 85 m.mu.Lock() 86 m.rEMA = rEMA 87 m.samples++ 88 m.mu.Unlock() 89 } 90 91 // IO is a convenience method intended to wrap io.Reader and io.Writer method 92 // execution. It calls m.Update(n) and then returns (n, err) unmodified. 93 func (m *Monitor) IO(n int, err error) (int, error) { 94 return m.Update(n), err 95 } 96 97 // Done marks the transfer as finished and prevents any further updates or 98 // limiting. Instantaneous and current transfer rates drop to 0. Update, IO, and 99 // Limit methods become NOOPs. It returns the total number of bytes transferred. 100 func (m *Monitor) Done() int64 { 101 m.mu.Lock() 102 if now := m.update(0); m.sBytes > 0 { 103 m.reset(now) 104 } 105 m.active = false 106 m.tLast = 0 107 n := m.bytes 108 m.mu.Unlock() 109 return n 110 } 111 112 // timeRemLimit is the maximum Status.TimeRem value. 113 const timeRemLimit = 999*time.Hour + 59*time.Minute + 59*time.Second 114 115 // Status represents the current Monitor status. All transfer rates are in bytes 116 // per second rounded to the nearest byte. 117 type Status struct { 118 Start time.Time // Transfer start time 119 Bytes int64 // Total number of bytes transferred 120 Samples int64 // Total number of samples taken 121 InstRate int64 // Instantaneous transfer rate 122 CurRate int64 // Current transfer rate (EMA of InstRate) 123 AvgRate int64 // Average transfer rate (Bytes / Duration) 124 PeakRate int64 // Maximum instantaneous transfer rate 125 BytesRem int64 // Number of bytes remaining in the transfer 126 Duration time.Duration // Time period covered by the statistics 127 Idle time.Duration // Time since the last transfer of at least 1 byte 128 TimeRem time.Duration // Estimated time to completion 129 Progress Percent // Overall transfer progress 130 Active bool // Flag indicating an active transfer 131 } 132 133 // Status returns current transfer status information. The returned value 134 // becomes static after a call to Done. 135 func (m *Monitor) Status() Status { 136 m.mu.Lock() 137 now := m.update(0) 138 s := Status{ 139 Active: m.active, 140 Start: clockToTime(m.start), 141 Duration: m.sLast - m.start, 142 Idle: now - m.tLast, 143 Bytes: m.bytes, 144 Samples: m.samples, 145 PeakRate: round(m.rPeak), 146 BytesRem: m.tBytes - m.bytes, 147 Progress: percentOf(float64(m.bytes), float64(m.tBytes)), 148 } 149 if s.BytesRem < 0 { 150 s.BytesRem = 0 151 } 152 if s.Duration > 0 { 153 rAvg := float64(s.Bytes) / s.Duration.Seconds() 154 s.AvgRate = round(rAvg) 155 if s.Active { 156 s.InstRate = round(m.rSample) 157 s.CurRate = round(m.rEMA) 158 if s.BytesRem > 0 { 159 if tRate := 0.8*m.rEMA + 0.2*rAvg; tRate > 0 { 160 ns := float64(s.BytesRem) / tRate * 1e9 161 if ns > float64(timeRemLimit) { 162 ns = float64(timeRemLimit) 163 } 164 s.TimeRem = clockRound(time.Duration(ns)) 165 } 166 } 167 } 168 } 169 m.mu.Unlock() 170 return s 171 } 172 173 // Limit restricts the instantaneous (per-sample) data flow to rate bytes per 174 // second. It returns the maximum number of bytes (0 <= n <= want) that may be 175 // transferred immediately without exceeding the limit. If block == true, the 176 // call blocks until n > 0. want is returned unmodified if want < 1, rate < 1, 177 // or the transfer is inactive (after a call to Done). 178 // 179 // At least one byte is always allowed to be transferred in any given sampling 180 // period. Thus, if the sampling rate is 100ms, the lowest achievable flow rate 181 // is 10 bytes per second. 182 // 183 // For usage examples, see the implementation of Reader and Writer in io.go. 184 func (m *Monitor) Limit(want int, rate int64, block bool) (n int) { 185 if want < 1 || rate < 1 { 186 return want 187 } 188 m.mu.Lock() 189 190 // Determine the maximum number of bytes that can be sent in one sample 191 limit := round(float64(rate) * m.sRate.Seconds()) 192 if limit <= 0 { 193 limit = 1 194 } 195 196 // If block == true, wait until m.sBytes < limit 197 if now := m.update(0); block { 198 for m.sBytes >= limit && m.active { 199 now = m.waitNextSample(now) 200 } 201 } 202 203 // Make limit <= want (unlimited if the transfer is no longer active) 204 if limit -= m.sBytes; limit > int64(want) || !m.active { 205 limit = int64(want) 206 } 207 m.mu.Unlock() 208 209 if limit < 0 { 210 limit = 0 211 } 212 return int(limit) 213 } 214 215 // SetTransferSize specifies the total size of the data transfer, which allows 216 // the Monitor to calculate the overall progress and time to completion. 217 func (m *Monitor) SetTransferSize(bytes int64) { 218 if bytes < 0 { 219 bytes = 0 220 } 221 m.mu.Lock() 222 m.tBytes = bytes 223 m.mu.Unlock() 224 } 225 226 // update accumulates the transferred byte count for the current sample until 227 // clock() - m.sLast >= m.sRate. The monitor status is updated once the current 228 // sample is done. 229 func (m *Monitor) update(n int) (now time.Duration) { 230 if !m.active { 231 return 232 } 233 if now = clock(); n > 0 { 234 m.tLast = now 235 } 236 m.sBytes += int64(n) 237 if sTime := now - m.sLast; sTime >= m.sRate { 238 t := sTime.Seconds() 239 if m.rSample = float64(m.sBytes) / t; m.rSample > m.rPeak { 240 m.rPeak = m.rSample 241 } 242 243 // Exponential moving average using a method similar to *nix load 244 // average calculation. Longer sampling periods carry greater weight. 245 if m.samples > 0 { 246 w := math.Exp(-t / m.rWindow) 247 m.rEMA = m.rSample + w*(m.rEMA-m.rSample) 248 } else { 249 m.rEMA = m.rSample 250 } 251 m.reset(now) 252 } 253 return 254 } 255 256 // reset clears the current sample state in preparation for the next sample. 257 func (m *Monitor) reset(sampleTime time.Duration) { 258 m.bytes += m.sBytes 259 m.samples++ 260 m.sBytes = 0 261 m.sLast = sampleTime 262 } 263 264 // waitNextSample sleeps for the remainder of the current sample. The lock is 265 // released and reacquired during the actual sleep period, so it's possible for 266 // the transfer to be inactive when this method returns. 267 func (m *Monitor) waitNextSample(now time.Duration) time.Duration { 268 const minWait = 5 * time.Millisecond 269 current := m.sLast 270 271 // sleep until the last sample time changes (ideally, just one iteration) 272 for m.sLast == current && m.active { 273 d := current + m.sRate - now 274 m.mu.Unlock() 275 if d < minWait { 276 d = minWait 277 } 278 time.Sleep(d) 279 m.mu.Lock() 280 now = m.update(0) 281 } 282 return now 283 }