bps/bandwidth.go

125 lines
2.7 KiB
Go
Raw Normal View History

// Package bps is a simple tool for keeping track of the rate of bytes
// transmitted
package bps
2014-03-02 22:44:09 -08:00
import (
2015-09-30 22:57:14 -07:00
"fmt"
"io/ioutil"
"sync"
2014-03-02 22:44:09 -08:00
"time"
)
2015-09-30 22:57:14 -07:00
const (
minResolution = 3
)
// BPS keeps track of state for byte counts
//
// Instantiate a BPS then feed bytes either via BPS.Add, or writing to it. When
// the accumulated values are needed call BPS.Cur. When the BPS is no loner
// needed call BPS.Close
type BPS struct {
sync.RWMutex
2014-03-02 23:03:00 -08:00
quit chan interface{}
closed chan interface{}
2014-03-02 23:03:00 -08:00
2015-09-30 22:57:14 -07:00
interval time.Duration
dt time.Duration
// curBs bytes read for this dt
curBs int64
2015-09-30 22:57:14 -07:00
// buckets contains an entry for bytes read for each dt of time in interval
buckets []int64
2015-09-30 22:57:14 -07:00
// timeI keys into buckets for the current point in time
timeI int
2014-03-02 23:03:00 -08:00
}
2015-09-30 22:57:14 -07:00
// New Returns a populated and ready to use BPS. interval is the amount of time
// (for example 60 seconds) over which to track byte flow (bytes/second for the
// last interval), and resolution is used in the following calculation:
//
// dt = interval / resolution (s)
//
// where the dt is the temporal resolution of the updates (add or remove
// information every dt).
func New(interval time.Duration, resolution uint) (*BPS, error) {
if resolution < minResolution {
return nil, fmt.Errorf("resolution must be larger than %d", minResolution)
}
2015-09-30 22:57:14 -07:00
dtns := interval.Nanoseconds() / int64(resolution)
dt := time.Duration(dtns) * time.Nanosecond
r := &BPS{
2015-09-30 22:57:14 -07:00
interval: interval,
dt: dt,
quit: make(chan interface{}),
closed: make(chan interface{}),
buckets: make([]int64, resolution),
2014-03-02 23:03:00 -08:00
}
go r.run()
2014-03-02 23:03:00 -08:00
return r, nil
}
func (b *BPS) runLoop() {
t := time.NewTicker(b.dt)
2014-03-02 23:03:00 -08:00
for {
select {
case <-t.C:
b.Lock()
2015-09-30 22:57:14 -07:00
b.buckets[b.timeI] = b.curBs
b.curBs = 0
2015-09-30 22:57:14 -07:00
b.timeI = (b.timeI + 1) % int(len(b.buckets))
b.Unlock()
case <-b.quit:
return
2014-03-02 23:03:00 -08:00
}
}
}
// Run is a method of BPS that must be started in a goroutine in order
// for things to be functional.
func (b *BPS) run() {
b.runLoop()
close(b.closed)
}
2015-09-30 11:04:39 -07:00
// Write implements io.Writer so that one can simply write bytes to the struct.
func (b *BPS) Write(p []byte) (int, error) {
n, err := ioutil.Discard.Write(p)
b.Add(int64(n))
return n, err
}
2015-09-30 11:04:39 -07:00
// Add adds i to the current time bucket.
func (b *BPS) Add(i int64) {
b.Lock()
b.curBs += i
2015-09-30 22:57:14 -07:00
b.buckets[b.timeI] = b.curBs
b.Unlock()
}
2015-09-30 22:57:14 -07:00
// Rate returns the current rate (bytes / second).
func (b *BPS) Rate() float64 {
b.Lock()
2015-09-30 22:57:14 -07:00
var total int64 = 0
for _, b := range b.buckets {
total += b
}
b.Unlock()
2015-09-30 22:57:14 -07:00
return float64(total) / b.interval.Seconds()
}
2015-09-30 23:07:41 -07:00
// HumanRate returns a human-friendly (e.g. 23.3MB/s) rate.
func (b *BPS) HumanRate() string {
return Bytes(uint64(b.Rate())) + "/s"
}
2015-09-30 11:04:39 -07:00
// Close cleans up and shuts down a BPS.
func (b *BPS) Close() {
close(b.quit)
<-b.closed
}