metrics: allow getting metrics by peer/protocol (#5)

metrics/bandwidth.go

@@ -94,3 +94,60 @@ func (bwc *BandwidthCounter) GetBandwidthTotals() (out Stats) {
 		RateOut:  outSnap.Rate,
 	}
 }
+
+// GetBandwidthByPeer returns a map of all remembered peers and the bandwidth
+// metrics with respect to each. This method may be very expensive.
+func (bwc *BandwidthCounter) GetBandwidthByPeer() map[peer.ID]Stats {
+	peers := make(map[peer.ID]Stats)
+
+	bwc.peerIn.ForEach(func(p string, meter *flow.Meter) {
+		id := peer.ID(p)
+		snap := meter.Snapshot()
+
+		stat := peers[id]
+		stat.TotalIn = int64(snap.Total)
+		stat.RateIn = snap.Rate
+		peers[id] = stat
+	})
+
+	bwc.peerOut.ForEach(func(p string, meter *flow.Meter) {
+		id := peer.ID(p)
+		snap := meter.Snapshot()
+
+		stat := peers[id]
+		stat.TotalOut = int64(snap.Total)
+		stat.RateOut = snap.Rate
+		peers[id] = stat
+	})
+
+	return peers
+}
+
+// GetBandwidthByProtocol returns a map of all remembered protocols and
+// the bandwidth metrics with respect to each. This method may be moderately
+// expensive.
+func (bwc *BandwidthCounter) GetBandwidthByProtocol() map[protocol.ID]Stats {
+	protocols := make(map[protocol.ID]Stats)
+
+	bwc.protocolIn.ForEach(func(p string, meter *flow.Meter) {
+		id := protocol.ID(p)
+		snap := meter.Snapshot()
+
+		stat := protocols[id]
+		stat.TotalIn = int64(snap.Total)
+		stat.RateIn = snap.Rate
+		protocols[id] = stat
+	})
+
+	bwc.protocolOut.ForEach(func(p string, meter *flow.Meter) {
+		id := protocol.ID(p)
+		snap := meter.Snapshot()
+
+		stat := protocols[id]
+		stat.TotalOut = int64(snap.Total)
+		stat.RateOut = snap.Rate
+		protocols[id] = stat
+	})
+
+	return protocols
+}

metrics/bandwidth_test.go

@@ -78,20 +78,44 @@ func TestBandwidthCounter(t *testing.T) {
 		}
 	}
 
+	assertProtocols := func(check func(Stats)) {
+		byProtocol := bwc.GetBandwidthByProtocol()
+		if len(byProtocol) != 2 {
+			t.Errorf("expected 2 protocols, got %d", len(byProtocol))
+		}
+		for i := 0; i < 2; i++ {
+			p := protocol.ID(fmt.Sprintf("proto-%d", i))
+			for _, stats := range [...]Stats{bwc.GetBandwidthForProtocol(p), byProtocol[p]} {
+				check(stats)
+			}
+		}
+	}
+
+	assertPeers := func(check func(Stats)) {
+		byPeer := bwc.GetBandwidthByPeer()
+		if len(byPeer) != 100 {
+			t.Errorf("expected 100 peers, got %d", len(byPeer))
+		}
+		for i := 0; i < 100; i++ {
+			p := peer.ID(fmt.Sprintf("peer-%d", i))
+			for _, stats := range [...]Stats{bwc.GetBandwidthForPeer(p), byPeer[p]} {
+				check(stats)
+			}
+		}
+	}
+
 	close(start)
 	time.Sleep(2*time.Second + 100*time.Millisecond)
 
-	for i := 0; i < 100; i++ {
-		stats := bwc.GetBandwidthForPeer(peer.ID(fmt.Sprintf("peer-%d", i)))
+	assertPeers(func(stats Stats) {
 		assertApproxEq(t, 2000, stats.RateOut)
 		assertApproxEq(t, 1000, stats.RateIn)
-	}
+	})
 
-	for i := 0; i < 2; i++ {
-		stats := bwc.GetBandwidthForProtocol(protocol.ID(fmt.Sprintf("proto-%d", i)))
+	assertProtocols(func(stats Stats) {
 		assertApproxEq(t, 100000, stats.RateOut)
 		assertApproxEq(t, 50000, stats.RateIn)
-	}
+	})
 
 	{
 		stats := bwc.GetBandwidthTotals()
@@ -101,17 +125,16 @@ func TestBandwidthCounter(t *testing.T) {
 
 	wg.Wait()
 	time.Sleep(1 * time.Second)
-	for i := 0; i < 100; i++ {
-		stats := bwc.GetBandwidthForPeer(peer.ID(fmt.Sprintf("peer-%d", i)))
+
+	assertPeers(func(stats Stats) {
 		assertEq(t, 8000, stats.TotalOut)
 		assertEq(t, 4000, stats.TotalIn)
-	}
+	})
 
-	for i := 0; i < 2; i++ {
-		stats := bwc.GetBandwidthForProtocol(protocol.ID(fmt.Sprintf("proto-%d", i)))
+	assertProtocols(func(stats Stats) {
 		assertEq(t, 400000, stats.TotalOut)
 		assertEq(t, 200000, stats.TotalIn)
-	}
+	})
 
 	{
 		stats := bwc.GetBandwidthTotals()
@@ -127,6 +150,7 @@ func assertEq(t *testing.T, expected, actual int64) {
 }
 
 func assertApproxEq(t *testing.T, expected, actual float64) {
+	t.Helper()
 	margin := expected * acceptableError
 	if !(math.Abs(expected-actual) <= margin) {
 		t.Errorf("expected %f (±%f), got %f", expected, margin, actual)

metrics/reporter.go

@@ -26,4 +26,6 @@ type Reporter interface {
 	GetBandwidthForPeer(peer.ID) Stats
 	GetBandwidthForProtocol(protocol.ID) Stats
 	GetBandwidthTotals() Stats
+	GetBandwidthByPeer() map[peer.ID]Stats
+	GetBandwidthByProtocol() map[protocol.ID]Stats
 }