refactor: implement TargetPinger.Statistics() as an iterator (consume…

…s/allocates slightly less memory)
clambin · Oct 17, 2024 · 423afb5 · 423afb5
1 parent a801789
commit 423afb5
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 21 deletions.
diff --git a/internal/cmd/pinger.go b/internal/cmd/pinger.go
@@ -7,7 +7,7 @@ import (
 	"github.com/clambin/pinger/internal/collector"
 	"github.com/clambin/pinger/internal/configuration"
 	"github.com/clambin/pinger/internal/pinger"
-	icmp2 "github.com/clambin/pinger/pkg/ping/icmp"
+	"github.com/clambin/pinger/pkg/ping/icmp"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/spf13/cobra"
@@ -33,12 +33,12 @@ var (
 func Main(cmd *cobra.Command, args []string) error {
 	l := charmer.GetLogger(cmd)
 	targets := configuration.GetTargets(viper.GetViper(), args)
-	var tp icmp2.Transport
+	var tp icmp.Transport
 	if viper.GetBool("ipv4") {
-		tp |= icmp2.IPv4
+		tp |= icmp.IPv4
 	}
 	if viper.GetBool("ipv6") {
-		tp |= icmp2.IPv6
+		tp |= icmp.IPv6
 	}
 
 	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)

diff --git a/internal/collector/collector.go b/internal/collector/collector.go
@@ -3,6 +3,7 @@ package collector
 import (
 	"github.com/clambin/pinger/pkg/ping"
 	"github.com/prometheus/client_golang/prometheus"
+	"iter"
 	"log/slog"
 )
 
@@ -34,7 +35,7 @@ type Collector struct {
 }
 
 type Pinger interface {
-	Statistics() map[string]ping.Statistics
+	Statistics() iter.Seq2[string, ping.Statistics]
 }
 
 // Describe implements the Prometheus Collector interface

diff --git a/internal/collector/collector_test.go b/internal/collector/collector_test.go
@@ -5,6 +5,7 @@ import (
 	"github.com/clambin/pinger/pkg/ping"
 	"github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/stretchr/testify/require"
+	"iter"
 	"log/slog"
 	"testing"
 	"time"
@@ -33,12 +34,12 @@ var _ Pinger = fakeTracker{}
 
 type fakeTracker struct{}
 
-func (f fakeTracker) Statistics() map[string]ping.Statistics {
-	return map[string]ping.Statistics{
-		"localhost": {
+func (f fakeTracker) Statistics() iter.Seq2[string, ping.Statistics] {
+	return func(yield func(string, ping.Statistics) bool) {
+		yield("localhost", ping.Statistics{
 			Sent:     20,
 			Received: 10,
 			Latency:  200 * time.Millisecond,
-		},
+		})
 	}
 }
diff --git a/internal/pinger/pinger.go b/internal/pinger/pinger.go
@@ -5,6 +5,7 @@ import (
 	"github.com/clambin/pinger/pkg/ping"
 	"github.com/clambin/pinger/pkg/ping/icmp"
 	"golang.org/x/exp/maps"
+	"iter"
 	"log/slog"
 	"time"
 )
@@ -48,11 +49,14 @@ func (tp *TargetPinger) Run(ctx context.Context) {
 	<-ctx.Done()
 }
 
-func (tp *TargetPinger) Statistics() map[string]ping.Statistics {
-	stats := make(map[string]ping.Statistics, len(tp.targets))
-	for name, target := range tp.targets {
-		stats[name] = target.Statistics()
-		target.ResetStatistics()
+func (tp *TargetPinger) Statistics() iter.Seq2[string, ping.Statistics] {
+	return func(yield func(string, ping.Statistics) bool) {
+		for name, target := range tp.targets {
+			stats := target.Statistics()
+			target.ResetStatistics()
+			if !yield(name, stats) {
+				return
+			}
+		}
 	}
-	return stats
 }
diff --git a/internal/pinger/pinger_test.go b/internal/pinger/pinger_test.go
@@ -38,12 +38,10 @@ func TestPinger(t *testing.T) {
 		return s.received.Load() > 1
 	}, 5*time.Second, time.Second)
 
-	stats := p.Statistics()
-	//t.Log(stats)
-	ipv4Stats, ok := stats["127.0.0.1"]
-	assert.True(t, ok)
-	assert.NotZero(t, ipv4Stats.Received)
-
+	for name, stats := range p.Statistics() {
+		assert.Equal(t, "127.0.0.1", name)
+		assert.NotZero(t, stats.Received)
+	}
 	cancel()
 	<-ch
 }