diff --git a/wgengine/bench/bench.go b/wgengine/bench/bench.go index 9c7fdd690..311a99697 100644 --- a/wgengine/bench/bench.go +++ b/wgengine/bench/bench.go @@ -80,7 +80,7 @@ func main() { // tx=134236 rx=133166 (1070 = 0.80% loss) (1088.9 Mbits/sec) case 101: - setupWGTest(logf, traf, Addr1, Addr2) + setupWGTest(nil, logf, traf, Addr1, Addr2) default: log.Fatalf("provide a valid test number (0..n)") diff --git a/wgengine/bench/bench_test.go b/wgengine/bench/bench_test.go index 7f8089445..253209bb6 100644 --- a/wgengine/bench/bench_test.go +++ b/wgengine/bench/bench_test.go @@ -43,7 +43,7 @@ func BenchmarkBatchTCP(b *testing.B) { func BenchmarkWireGuardTest(b *testing.B) { run(b, func(logf logger.Logf, traf *TrafficGen) { - setupWGTest(logf, traf, Addr1, Addr2) + setupWGTest(b, logf, traf, Addr1, Addr2) }) } diff --git a/wgengine/bench/wg.go b/wgengine/bench/wg.go index 167975a3f..e3cb300f4 100644 --- a/wgengine/bench/wg.go +++ b/wgengine/bench/wg.go @@ -10,6 +10,7 @@ import ( "os" "strings" "sync" + "testing" "github.com/tailscale/wireguard-go/tun" "inet.af/netaddr" @@ -25,7 +26,7 @@ import ( "tailscale.com/wgengine/wgcfg" ) -func setupWGTest(logf logger.Logf, traf *TrafficGen, a1, a2 netaddr.IPPrefix) { +func setupWGTest(b *testing.B, logf logger.Logf, traf *TrafficGen, a1, a2 netaddr.IPPrefix) { l1 := logger.WithPrefix(logf, "e1: ") k1, err := wgkey.NewPrivate() if err != nil { @@ -49,6 +50,9 @@ func setupWGTest(logf logger.Logf, traf *TrafficGen, a1, a2 netaddr.IPPrefix) { if err != nil { log.Fatalf("e1 init: %v", err) } + if b != nil { + b.Cleanup(e1.Close) + } l2 := logger.WithPrefix(logf, "e2: ") k2, err := wgkey.NewPrivate() @@ -73,6 +77,9 @@ func setupWGTest(logf logger.Logf, traf *TrafficGen, a1, a2 netaddr.IPPrefix) { if err != nil { log.Fatalf("e2 init: %v", err) } + if b != nil { + b.Cleanup(e2.Close) + } e1.SetFilter(filter.NewAllowAllForTest(l1)) e2.SetFilter(filter.NewAllowAllForTest(l2))