fix(agent): prevent goroutine pile up in reportMetadataLoop

ammario · ammario · commit bb1ef419e149 · 2023-04-01T18:43:01.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,11 +35,11 @@ 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"
 	"tailscale.com/types/netlogtype"
+	"tailscale.com/util/singleflight"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/agent/usershell"
@@ -276,7 +276,7 @@ func (a *agent) reportMetadataLoop(ctx context.Context) {
 	)
 	defer baseTicker.Stop()
 
-	var flight singleflight.Group
+	var flight singleflight.Group[string, struct{}]
 
 	for {
 		select {
@@ -345,10 +345,11 @@ func (a *agent) reportMetadataLoop(ctx context.Context) {
 			}
 
 			md := md
+			lastCollectedAts[md.Key] = time.Now()
 			// 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) {
+			flight.DoChan(md.Key, func() (struct{}, error) {
 				timeout := md.Timeout
 				if timeout == 0 {
 					timeout = md.Interval
@@ -360,13 +361,12 @@ 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
+				return struct{}{}, nil
 			})
 		}
 	}
