github.com/ethersphere/bee/v2@v2.2.0/pkg/rate/rate.go (about)

     1  // Copyright 2022 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package rate
     6  
     7  // Package rate is a thread-safe rate tracker with per second resolution.
     8  // Under the hood, it uses a moving window to interpolate a rate.
     9  
    10  import (
    11  	"sync"
    12  	"time"
    13  )
    14  
    15  const milliInSeconds = 1000
    16  
    17  type Rate struct {
    18  	mtx        sync.Mutex
    19  	windows    map[int64]int
    20  	windowSize int64 // window size in milliseconds
    21  	now        func() time.Time
    22  }
    23  
    24  // New returns a new rate tracker with a defined window size that must be greater than one millisecond.
    25  func New(windowsSize time.Duration) *Rate {
    26  	return &Rate{
    27  		windows:    make(map[int64]int),
    28  		windowSize: windowsSize.Milliseconds(),
    29  		now:        func() time.Time { return time.Now() },
    30  	}
    31  
    32  }
    33  
    34  // add uses the current time and rounds it down to a window
    35  // and increments the window's value.
    36  func (r *Rate) Add(count int) {
    37  	r.mtx.Lock()
    38  	defer r.mtx.Unlock()
    39  	defer r.cleanup()
    40  
    41  	window := r.now().UnixMilli() / r.windowSize
    42  	r.windows[window] += count
    43  }
    44  
    45  // rate uses the current window and previous windows' counter to compute a moving-window rate in seconds.
    46  // the rate is computed by first calculating how far along the current time is in the current window
    47  // as a ratio between 0 and 1.0. Then, the sum of currentWindowCount + (1 - ratio) * previousWindowCounter
    48  // is returned as the interpolated counter between the two windows.
    49  func (r *Rate) Rate() float64 {
    50  	r.mtx.Lock()
    51  	defer r.mtx.Unlock()
    52  	defer r.cleanup()
    53  
    54  	now := r.now().UnixMilli()
    55  	window := now / r.windowSize
    56  
    57  	interpolate := 1 - float64(now-(window*r.windowSize))/float64(r.windowSize)
    58  
    59  	return milliInSeconds * ((float64(r.windows[window])) + ((interpolate) * float64(r.windows[window-1]))) / float64(r.windowSize)
    60  }
    61  
    62  // cleanup removes windows older than the most recent two windows.
    63  // Must be called under lock.
    64  func (r *Rate) cleanup() {
    65  
    66  	window := r.now().UnixMilli() / r.windowSize
    67  	for k := range r.windows {
    68  		if k <= window-2 {
    69  			delete(r.windows, k)
    70  		}
    71  	}
    72  }