From 996387510d8b366a1db349126df6eafb22d0ee10 Mon Sep 17 00:00:00 2001
From: Cole Brown <bigswim@gmail.com>
Date: Wed, 27 Jun 2018 13:25:16 -0400
Subject: [PATCH] Run benchmarks 10s, report rates

---
 benchmark_utils.go | 17 +++++++++++------
 peerstore_test.go  | 46 ++++++++++++++++++++++++++++++++++------------
 2 files changed, 45 insertions(+), 18 deletions(-)

diff --git a/benchmark_utils.go b/benchmark_utils.go
index c66fd07..7921090 100644
--- a/benchmark_utils.go
+++ b/benchmark_utils.go
@@ -45,18 +45,23 @@ func randomPeer(b *testing.B) *peerpair {
 	return &peerpair{b58ID, addr}
 }
 
-func addressProducer(b *testing.B, addrs chan *peerpair, n int) {
+func addressProducer(b *testing.B, addrs chan *peerpair, done <-chan time.Time) {
 	defer close(addrs)
-	for i := 0; i < n; i++ {
-		p := randomPeer(b)
-		addrs <- p
+	for {
+		select {
+		case <-done:
+			return
+		default:
+			p := randomPeer(b)
+			addrs <- p
+		}
 	}
 }
 
-func rateLimitedAddressProducer(b *testing.B, addrs chan *peerpair, producer chan *peerpair, n int, avgTime time.Duration, errBound time.Duration) {
+func rateLimitedAddressProducer(b *testing.B, addrs chan *peerpair, producer chan *peerpair, done <-chan time.Time, avgTime time.Duration, errBound time.Duration) {
 	defer close(addrs)
 
-	go addressProducer(b, producer, n)
+	go addressProducer(b, producer, done)
 
 	eb := int64(errBound)
 	for {
diff --git a/peerstore_test.go b/peerstore_test.go
index e6d9407..0efbf4c 100644
--- a/peerstore_test.go
+++ b/peerstore_test.go
@@ -309,39 +309,61 @@ func TestBasicPeerstore(t *testing.T) {
 	runTestWithPeerstores(t, testBasicPeerstore)
 }
 
-func BenchmarkBasicPeerstore(b *testing.B) {
-	ps := NewPeerstore()
+func benchmarkPeerstore(ps Peerstore) func(*testing.B) {
+	return func(b *testing.B) {
+		addrs := make(chan *peerpair, 100)
 
-	addrs := make(chan *peerpair, 100)
+		go addressProducer(b, addrs, time.After(10*time.Second))
 
-	go addressProducer(b, addrs, 50000)
-
-	for {
-		pp, ok := <-addrs
-		if !ok {
-			break
+		start := time.Now()
+		count := 0
+		b.ResetTimer()
+		for {
+			pp, ok := <-addrs
+			if !ok {
+				break
+			}
+			count++
+			pid := peer.ID(pp.ID)
+			ps.AddAddr(pid, pp.Addr, PermanentAddrTTL)
 		}
-		pid := peer.ID(pp.ID)
-		ps.AddAddr(pid, pp.Addr, PermanentAddrTTL)
+		elapsed := time.Since(start)
+		rate := float64(count) / elapsed.Seconds()
+		b.Logf("Added %d addresses in %s, rate: %f addrs/s\n", count, elapsed, rate)
 	}
 }
 
+func BenchmarkPeerstore(b *testing.B) {
+	ps := NewPeerstore()
+	b.Run("PeerstoreBasic", benchmarkPeerstore(ps))
+
+	dsps, closer := setupDatastorePeerstore(b)
+	defer closer()
+	b.Run("PeerstoreDatastore", benchmarkPeerstore(dsps))
+}
+
 func benchmarkPeerstoreRateLimited(ps Peerstore) func(*testing.B) {
 	return func(b *testing.B) {
 		producer := make(chan *peerpair, 100)
 		addrs := make(chan *peerpair, 100)
 
-		go rateLimitedAddressProducer(b, addrs, producer, 60000, 100*time.Microsecond, 100*time.Microsecond)
+		go rateLimitedAddressProducer(b, addrs, producer, time.After(10*time.Second), 100*time.Microsecond, 100*time.Microsecond)
 
+		start := time.Now()
+		count := 0
 		b.ResetTimer()
 		for {
 			pp, ok := <-addrs
 			if !ok {
 				break
 			}
+			count++
 			pid := peer.ID(pp.ID)
 			ps.AddAddr(pid, pp.Addr, PermanentAddrTTL)
 		}
+		elapsed := time.Since(start)
+		rate := float64(count) / elapsed.Seconds()
+		b.Logf("Added %d addresses in %s, rate: %f addrs/s\n", count, elapsed, rate)
 	}
 }