bps/bandwidth.go

// Package bps is a simple tool for keeping track of the rate of bytes
// transmitted
package bps

import (
	"errors"
	"io/ioutil"
	"sort"
	"sync"
	"time"
)

// 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

	quit   chan interface{}
	closed chan interface{}

	snapshot []float64

	dt  time.Duration
	dts []int64

	// curBs bytes read for this dt
	curBs int64

	// timeBuckets contains an entry for bytes read for each dt of time up to
	// the longest recoreded time slice.
	timeBuckets []int64

	// timeI keys into timeBuckets for the current point in time
	timeI int

	// max is defined in New to be the maximum number of temporal buckets
	// required.
	max int64
}

// New Returns a populated and ready to launch BPS. dts is
// a slice of multiples of dt on which to report (e.g. 1x, 10x, 60x dt). dt is
// also how often the values used to send to Rx and Tx are updated.
func New(dts []int, dt time.Duration) (*BPS, error) {
	if len(dts) < 1 {
		return nil, errors.New("must specify at least one interval lenght")
	}
	sort.Ints(dts)
	convertedDts := []int64{}
	for _, dt := range dts {
		convertedDts = append(convertedDts, int64(dt))
	}
	max := convertedDts[len(convertedDts)-1]
	r := &BPS{
		dt:          dt,
		dts:         convertedDts,
		quit:        make(chan interface{}),
		closed:      make(chan interface{}),
		timeBuckets: make([]int64, max),
		max:         max,
	}
	go r.run()
	return r, nil
}

func (b *BPS) runLoop() {
	t := time.NewTicker(b.dt)
	for {
		select {
		case <-t.C:
			b.Lock()

			b.timeBuckets[b.timeI] = b.curBs

			b.snapshot = b.averages(b.timeBuckets)

			b.curBs = 0

			// Here we march forward through time by going backward in our
			// slice.
			b.timeI = (b.timeI - 1) % int(b.max)
			// because modulo does unexpected things for negative numbers.
			if b.timeI < 0 {
				b.timeI = b.timeI + int(b.max)
			}
			b.Unlock()
		case <-b.quit:
			return
		}
	}
}

// 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)
}

func (b *BPS) averages(state []int64) []float64 {
	r := []float64{}
	var i int64 = 0
	var total int64 = 0
	for _, ti := range b.dts {
		for ; ; i++ {
			if i == ti {
				break
			}
			total += state[(int64(b.timeI)+i)%b.max]
		}
		r = append(r, float64(total)/float64(ti))
	}
	return r
}

func (b *BPS) Write(p []byte) (int, error) {
	n, err := ioutil.Discard.Write(p)
	b.Add(int64(n))
	return n, err
}

func (b *BPS) Add(i int64) {
	b.Lock()
	b.curBs += i
	b.timeBuckets[b.timeI] = b.curBs
	b.snapshot = b.averages(b.timeBuckets)
	b.Unlock()
}

func (b *BPS) Cur() []float64 {
	r := make([]float64, len(b.dts))
	b.Lock()
	for i := range b.snapshot {
		r[i] = b.snapshot[i]
	}
	b.Unlock()
	return r
}

func (b *BPS) Close() {
	close(b.quit)
	<-b.closed
}