fix(agent): prevent goroutine pile up in reportMetadataLoop

ammario · ammario · commit f82b5b6889df · 2023-04-01T18:50:24.000Z
In the prior implementation, calls to DoChan would stack up because we
weren't updating lastCollectedAts until collectMetadata finished. This
wasn't a true leak, instead, it meant that there would be
up to ~ (collectionRuntime / baseInterval) outstanding goroutines. So, for
example, if `sleep 60s` was the metadata script there would be up to
60 goroutines waiting at peak.
diff --git a/agent/agent.go b/agent/agent.go
@@ -35,7 +35,6 @@ import (
 	"go.uber.org/atomic"
 	gossh "golang.org/x/crypto/ssh"
 	"golang.org/x/exp/slices"
-	"golang.org/x/sync/singleflight"
 	"golang.org/x/xerrors"
 	"tailscale.com/net/speedtest"
 	"tailscale.com/tailcfg"
@@ -264,6 +263,21 @@ type metadataResultAndKey struct {
 	key    string
 }
 
+type trySingleflight struct {
+	m sync.Map
+}
+
+func (t *trySingleflight) Do(key string, fn func()) {
+	_, loaded := t.m.LoadOrStore(key, struct{}{})
+	if !loaded {
+		// There is already a goroutine running for this key.
+		return
+	}
+
+	defer t.m.Delete(key)
+	fn()
+}
+
 func (a *agent) reportMetadataLoop(ctx context.Context) {
 	baseInterval := adjustIntervalForTests(1)
 
@@ -276,7 +290,7 @@ func (a *agent) reportMetadataLoop(ctx context.Context) {
 	)
 	defer baseTicker.Stop()
 
-	var flight singleflight.Group
+	var flight trySingleflight
 
 	for {
 		select {
@@ -348,7 +362,7 @@ func (a *agent) reportMetadataLoop(ctx context.Context) {
 			// We send the result to the channel in the goroutine to avoid
 			// sending the same result multiple times. So, we don't care about
 			// the return values.
-			flight.DoChan(md.Key, func() (interface{}, error) {
+			go flight.Do(md.Key, func() {
 				timeout := md.Timeout
 				if timeout == 0 {
 					timeout = md.Interval
@@ -360,13 +374,11 @@ func (a *agent) reportMetadataLoop(ctx context.Context) {
 
 				select {
 				case <-ctx.Done():
-					return 0, nil
 				case metadataResults <- metadataResultAndKey{
 					key:    md.Key,
 					result: a.collectMetadata(ctx, md),
 				}:
 				}
-				return 0, nil
 			})
 		}
 	}
