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  }