From a72fb7ac0bd3d6be68fc019ead356941660333a6 Mon Sep 17 00:00:00 2001
From: Josh Bleecher Snyder <josh@tailscale.com>
Date: Fri, 7 May 2021 12:55:47 -0700
Subject: [PATCH] wgengine/bench: handle multiple Engine status callbacks

It is possible to get multiple status callbacks from an Engine.
We need to wait for at least one from each Engine.
Without limiting to one per Engine,
wait.Wait can exit early or can panic due to a negative counter.

Signed-off-by: Josh Bleecher Snyder <josh@tailscale.com>
---
 wgengine/bench/wg.go | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/wgengine/bench/wg.go b/wgengine/bench/wg.go
index e3cb300f4..412799e15 100644
--- a/wgengine/bench/wg.go
+++ b/wgengine/bench/wg.go
@@ -87,6 +87,7 @@ func setupWGTest(b *testing.B, logf logger.Logf, traf *TrafficGen, a1, a2 netadd
 	var wait sync.WaitGroup
 	wait.Add(2)
 
+	var e1waitDoneOnce sync.Once
 	e1.SetStatusCallback(func(st *wgengine.Status, err error) {
 		if err != nil {
 			log.Fatalf("e1 status err: %v", err)
@@ -118,9 +119,10 @@ func setupWGTest(b *testing.B, logf logger.Logf, traf *TrafficGen, a1, a2 netadd
 		}
 		c2.Peers = []wgcfg.Peer{p}
 		e2.Reconfig(&c2, &router.Config{}, new(dns.Config))
-		wait.Done()
+		e1waitDoneOnce.Do(wait.Done)
 	})
 
+	var e2waitDoneOnce sync.Once
 	e2.SetStatusCallback(func(st *wgengine.Status, err error) {
 		if err != nil {
 			log.Fatalf("e2 status err: %v", err)
@@ -152,7 +154,7 @@ func setupWGTest(b *testing.B, logf logger.Logf, traf *TrafficGen, a1, a2 netadd
 		}
 		c1.Peers = []wgcfg.Peer{p}
 		e1.Reconfig(&c1, &router.Config{}, new(dns.Config))
-		wait.Done()
+		e2waitDoneOnce.Do(wait.Done)
 	})
 
 	// Not using DERP in this test (for now?).