From 6a9cba0e1b17da2b986b32aa6e18e939fa02051a Mon Sep 17 00:00:00 2001
From: Cole Brown <bigswim@gmail.com>
Date: Tue, 26 Jun 2018 20:10:17 -0400
Subject: [PATCH] Update benchmarks to run against both peerstores

---
 addr_manager_test.go |  2 +-
 peerstore_test.go    | 37 +++++++++++++++++++++++--------------
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/addr_manager_test.go b/addr_manager_test.go
index 4adf84d..8498bc9 100644
--- a/addr_manager_test.go
+++ b/addr_manager_test.go
@@ -53,7 +53,7 @@ func testHas(t *testing.T, exp, act []ma.Multiaddr) {
 	}
 }
 
-func setupBadgerDatastore(t *testing.T) (datastore.Batching, func()) {
+func setupBadgerDatastore(t testing.TB) (datastore.Batching, func()) {
 	dataPath, err := ioutil.TempDir(os.TempDir(), "badger")
 	if err != nil {
 		t.Fatal(err)
diff --git a/peerstore_test.go b/peerstore_test.go
index 1799ce8..58b0619 100644
--- a/peerstore_test.go
+++ b/peerstore_test.go
@@ -38,7 +38,7 @@ func runTestWithPeerstores(t *testing.T, testFunc func(*testing.T, Peerstore)) {
 	testFunc(t, ps2)
 }
 
-func setupDatastorePeerstore(t *testing.T) (Peerstore, func()) {
+func setupDatastorePeerstore(t testing.TB) (Peerstore, func()) {
 	ds, closeDB := setupBadgerDatastore(t)
 	ctx, cancel := context.WithCancel(context.Background())
 	ps, err := NewPeerstoreDatastore(ctx, ds)
@@ -326,21 +326,30 @@ func BenchmarkBasicPeerstore(b *testing.B) {
 	}
 }
 
-func BenchmarkBasicPeerstoreRateLimited(b *testing.B) {
-	ps := NewPeerstore()
+func benchmarkPeerstoreRateLimited(ps Peerstore) func(*testing.B) {
+	return func(b *testing.B) {
+		producer := make(chan *peerpair, 100)
+		addrs := make(chan *peerpair, 100)
 
-	producer := make(chan *peerpair, 100)
-	addrs := make(chan *peerpair, 100)
+		go rateLimitedAddressProducer(b, addrs, producer, 10000, time.Millisecond, 500*time.Microsecond)
 
-	go rateLimitedAddressProducer(b, addrs, producer, 5000, time.Millisecond, 500*time.Microsecond)
-
-	b.ResetTimer()
-	for {
-		pp, ok := <-addrs
-		if !ok {
-			break
+		b.ResetTimer()
+		for {
+			pp, ok := <-addrs
+			if !ok {
+				break
+			}
+			pid := peer.ID(pp.ID)
+			ps.AddAddr(pid, pp.Addr, PermanentAddrTTL)
 		}
-		pid := peer.ID(pp.ID)
-		ps.AddAddr(pid, pp.Addr, PermanentAddrTTL)
 	}
 }
+
+func BenchmarkBasicPeerstoreRateLimited(b *testing.B) {
+	ps := NewPeerstore()
+	b.Run("PeerstoreBasic", benchmarkPeerstoreRateLimited(ps))
+
+	dsps, closer := setupDatastorePeerstore(b)
+	defer closer()
+	b.Run("PeerstoreDatastore", benchmarkPeerstoreRateLimited(dsps))
+}