Make rate estimator estimate packet rates too.

estimator/estimator.go

@@ -8,12 +8,14 @@ import (
 
 type Estimator struct {
 	interval time.Duration
-	count    uint32
+	bytes    uint32
+	packets  uint32
 
 	mu           sync.Mutex
 	totalBytes   uint32
 	totalPackets uint32
 	rate         uint32
+	packetRate   uint32
 	time         time.Time
 }
 
@@ -26,30 +28,38 @@ func New(interval time.Duration) *Estimator {
 
 func (e *Estimator) swap(now time.Time) {
 	interval := now.Sub(e.time)
-	count := atomic.SwapUint32(&e.count, 0)
+	bytes := atomic.SwapUint32(&e.bytes, 0)
+	packets := atomic.SwapUint32(&e.packets, 0)
+	atomic.AddUint32(&e.totalBytes, bytes)
+	atomic.AddUint32(&e.totalPackets, packets)
+
 	if interval < time.Millisecond {
 		e.rate = 0
+		e.packetRate = 0
 	} else {
-		e.rate = uint32(uint64(count*1000) / uint64(interval/time.Millisecond))
+		e.rate = uint32(uint64(bytes*1000) /
+			uint64(interval/time.Millisecond))
+		e.packetRate = uint32(uint64(packets*1000) /
+			uint64(interval/time.Millisecond))
+
 	}
 	e.time = now
 }
 
 func (e *Estimator) Accumulate(count uint32) {
-	atomic.AddUint32(&e.totalBytes, count)
-	atomic.AddUint32(&e.totalPackets, 1)
-	atomic.AddUint32(&e.count, count)
+	atomic.AddUint32(&e.bytes, count)
+	atomic.AddUint32(&e.packets, 1)
 }
 
-func (e *Estimator) estimate(now time.Time) uint32 {
+func (e *Estimator) estimate(now time.Time) (uint32, uint32) {
 	if now.Sub(e.time) > e.interval {
 		e.swap(now)
 	}
 
-	return e.rate
+	return e.rate, e.packetRate
 }
 
-func (e *Estimator) Estimate() uint32 {
+func (e *Estimator) Estimate() (uint32, uint32) {
 	now := time.Now()
 
 	e.mu.Lock()
@@ -58,7 +68,7 @@ func (e *Estimator) Estimate() uint32 {
 }
 
 func (e *Estimator) Totals() (uint32, uint32) {
-	b := atomic.LoadUint32(&e.totalBytes)
-	p := atomic.LoadUint32(&e.totalPackets)
+	b := atomic.LoadUint32(&e.totalBytes) + atomic.LoadUint32(&e.bytes)
+	p := atomic.LoadUint32(&e.totalPackets) + atomic.LoadUint32(&e.packets)
 	return p, b
 }

estimator/estimator_test.go

@@ -13,11 +13,14 @@ func TestEstimator(t *testing.T) {
 	e.Accumulate(42)
 	e.Accumulate(128)
 	e.estimate(now.Add(time.Second))
-	rate := e.estimate(now.Add(time.Second + time.Millisecond))
+	rate, packetRate := e.estimate(now.Add(time.Second + time.Millisecond))
 
 	if rate != 42+128 {
 		t.Errorf("Expected %v, got %v", 42+128, rate)
 	}
+	if packetRate != 2 {
+		t.Errorf("Expected 2, got %v", packetRate)
+	}
 
 	totalP, totalB := e.Totals()
 	if totalP != 2 {
@@ -26,4 +29,15 @@ func TestEstimator(t *testing.T) {
 	if totalB != 42+128 {
 		t.Errorf("Expected %v, got %v", 42+128, totalB)
 	}
+
+	e.Accumulate(12)
+
+	totalP, totalB = e.Totals()
+	if totalP != 3 {
+		t.Errorf("Expected 2, got %v", totalP)
+	}
+	if totalB != 42+128+12 {
+		t.Errorf("Expected %v, got %v", 42+128, totalB)
+	}
+
 }