From fecb67fed168c57d98372fca73f609a973aff162 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 8 Jul 2022 23:17:46 +0000
Subject: [PATCH] introduce an API to set a mock clock

---
 .travis.yml                      |  2 +-
 go.mod                           |  4 ++-
 go.sum                           |  2 ++
 mockclocktest/mock_clock_test.go | 43 ++++++++++++++++++++++++++++++++
 sweeper.go                       | 16 +++++++++---
 5 files changed, 62 insertions(+), 5 deletions(-)
 create mode 100644 go.sum
 create mode 100644 mockclocktest/mock_clock_test.go

diff --git a/.travis.yml b/.travis.yml
index a156d3e..67e608f 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -4,7 +4,7 @@ os:
 language: go
 
 go:
-  - 1.13.x
+  - 1.17.x
 
 env:
   global:
diff --git a/go.mod b/go.mod
index e9b50d4..0c94703 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,5 @@
 module github.com/libp2p/go-flow-metrics
 
-go 1.12
+go 1.17
+
+require github.com/benbjohnson/clock v1.3.0
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..c284fd6
--- /dev/null
+++ b/go.sum
@@ -0,0 +1,2 @@
+github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
+github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
diff --git a/mockclocktest/mock_clock_test.go b/mockclocktest/mock_clock_test.go
new file mode 100644
index 0000000..9cff3d6
--- /dev/null
+++ b/mockclocktest/mock_clock_test.go
@@ -0,0 +1,43 @@
+package mockclocktest
+
+import (
+	"testing"
+	"time"
+
+	"github.com/libp2p/go-flow-metrics"
+
+	"github.com/benbjohnson/clock"
+)
+
+var cl = clock.NewMock()
+
+func init() {
+	flow.SetClock(cl)
+}
+
+func TestBasic(t *testing.T) {
+	m := new(flow.Meter)
+	for i := 0; i < 300; i++ {
+		m.Mark(1000)
+		cl.Add(40 * time.Millisecond)
+	}
+	if rate := m.Snapshot().Rate; rate != 25000 {
+		t.Errorf("expected rate 25000, got %f", rate)
+	}
+
+	for i := 0; i < 200; i++ {
+		m.Mark(200)
+		cl.Add(40 * time.Millisecond)
+	}
+
+	// Adjusts
+	if rate := m.Snapshot().Rate; rate != 5017.776503840969 {
+		t.Errorf("expected rate 5017.776503840969, got %f", rate)
+	}
+
+	// Let it settle.
+	cl.Add(2 * time.Second)
+	if total := m.Snapshot().Total; total != 340000 {
+		t.Errorf("expected total 3400000, got %d", total)
+	}
+}
diff --git a/sweeper.go b/sweeper.go
index ceed13f..8b9a262 100644
--- a/sweeper.go
+++ b/sweeper.go
@@ -5,6 +5,8 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+
+	"github.com/benbjohnson/clock"
 )
 
 // IdleRate the rate at which we declare a meter idle (and stop tracking it
@@ -20,6 +22,14 @@ var alpha = 1 - math.Exp(-1.0)
 // The global sweeper.
 var globalSweeper sweeper
 
+var cl = clock.New()
+
+// SetClock sets a clock to use in the sweeper.
+// This will probably only ever be useful for testing purposes.
+func SetClock(c clock.Clock) {
+	cl = c
+}
+
 type sweeper struct {
 	sweepOnce sync.Once
 
@@ -53,10 +63,10 @@ func (sw *sweeper) register(m *Meter) {
 }
 
 func (sw *sweeper) runActive() {
-	ticker := time.NewTicker(time.Second)
+	ticker := cl.Ticker(time.Second)
 	defer ticker.Stop()
 
-	sw.lastUpdateTime = time.Now()
+	sw.lastUpdateTime = cl.Now()
 	for len(sw.meters) > 0 {
 		// Scale back allocation.
 		if len(sw.meters)*2 < cap(sw.meters) {
@@ -80,7 +90,7 @@ func (sw *sweeper) update() {
 	sw.snapshotMu.Lock()
 	defer sw.snapshotMu.Unlock()
 
-	now := time.Now()
+	now := cl.Now()
 	tdiff := now.Sub(sw.lastUpdateTime)
 	if tdiff <= 0 {
 		return