From a6db6b105d25e3198eee5c4006ea304704e4bcf7 Mon Sep 17 00:00:00 2001
From: Juliusz Chroboczek <jch@irif.fr>
Date: Sun, 31 May 2020 23:54:55 +0200
Subject: [PATCH] Keep track of total numbers sent.

---
 estimator/estimator.go      | 16 +++++++++++++---
 estimator/estimator_test.go |  8 ++++++++
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/estimator/estimator.go b/estimator/estimator.go
index 7ac8aa6..7b00f39 100644
--- a/estimator/estimator.go
+++ b/estimator/estimator.go
@@ -10,9 +10,11 @@ type Estimator struct {
 	interval time.Duration
 	count    uint32
 
-	mu   sync.Mutex
-	rate uint32
-	time time.Time
+	mu           sync.Mutex
+	totalBytes   uint32
+	totalPackets uint32
+	rate         uint32
+	time         time.Time
 }
 
 func New(interval time.Duration) *Estimator {
@@ -34,6 +36,8 @@ func (e *Estimator) swap(now time.Time) {
 }
 
 func (e *Estimator) Accumulate(count uint32) {
+	atomic.AddUint32(&e.totalBytes, count)
+	atomic.AddUint32(&e.totalPackets, 1)
 	atomic.AddUint32(&e.count, count)
 }
 
@@ -52,3 +56,9 @@ func (e *Estimator) Estimate() uint32 {
 	defer e.mu.Unlock()
 	return e.estimate(now)
 }
+
+func (e *Estimator) Totals() (uint32, uint32) {
+	b := atomic.LoadUint32(&e.totalBytes)
+	p := atomic.LoadUint32(&e.totalPackets)
+	return p, b
+}
diff --git a/estimator/estimator_test.go b/estimator/estimator_test.go
index 1fca1bc..ed5926e 100644
--- a/estimator/estimator_test.go
+++ b/estimator/estimator_test.go
@@ -18,4 +18,12 @@ func TestEstimator(t *testing.T) {
 	if rate != 42+128 {
 		t.Errorf("Expected %v, got %v", 42+128, rate)
 	}
+
+	totalP, totalB := e.Totals()
+	if totalP != 2 {
+		t.Errorf("Expected 2, got %v", totalP)
+	}
+	if totalB != 42+128 {
+		t.Errorf("Expected %v, got %v", 42+128, totalB)
+	}
 }