CloseFunc

coder · mtojek · Apr 14, 2023 · Apr 3, 2023 · Apr 4, 2023 · Apr 4, 2023
commit 9c7aef8064dbbcb7df2a79f76a67cad9dd1626d3
diff --git a/coderd/prometheusmetrics/prometheusmetrics.go b/coderd/prometheusmetrics/prometheusmetrics.go
@@ -293,7 +293,7 @@ func Agents(ctx context.Context, logger slog.Logger, registerer prometheus.Regis
 	return cancelFunc, nil
 }
 
-func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.Registerer, db database.Store, initialCreateAfter time.Time, duration time.Duration) (context.CancelFunc, error) {
+func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.Registerer, db database.Store, initialCreateAfter time.Time, duration time.Duration) (func(), error) {
 	if duration == 0 {
 		duration = 1 * time.Minute
 	}
@@ -398,10 +398,13 @@ func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.R
 		return nil, err
 	}
 
-	createdAfter := initialCreateAfter
 	ctx, cancelFunc := context.WithCancel(ctx)
+	done := make(chan struct{})
+
+	createdAfter := initialCreateAfter
 	ticker := time.NewTicker(duration)
 	go func() {
+		defer close(done)
 		defer ticker.Stop()
 		for {
 			select {
@@ -452,5 +455,8 @@ func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.R
 			ticker.Reset(duration)
 		}
 	}()
-	return cancelFunc, nil
+	return func() {
+		cancelFunc()
+		<-done
+	}, nil
 }
diff --git a/coderd/prometheusmetrics/prometheusmetrics_test.go b/coderd/prometheusmetrics/prometheusmetrics_test.go
@@ -367,9 +367,9 @@ func TestAgentStats(t *testing.T) {
 
 	user := coderdtest.CreateFirstUser(t, client)
 
-	agent1, _ := prepareWorkspaceAndAgent(t, client, user, 1)
-	agent2, _ := prepareWorkspaceAndAgent(t, client, user, 2)
-	agent3, _ := prepareWorkspaceAndAgent(t, client, user, 3)
+	agent1 := prepareWorkspaceAndAgent(t, client, user, 1)
+	agent2 := prepareWorkspaceAndAgent(t, client, user, 2)
+	agent3 := prepareWorkspaceAndAgent(t, client, user, 3)
 
 	registry := prometheus.NewRegistry()
 
@@ -409,9 +409,9 @@ func TestAgentStats(t *testing.T) {
 	//
 	// Set initialCreateAfter to some time in the past, so that AgentStats would include all above PostStats,
 	// and it doesn't depend on the real time.
-	cancel, err := prometheusmetrics.AgentStats(ctx, slogtest.Make(t, nil), registry, db, time.Now().Add(-time.Minute), time.Millisecond)
+	closeFunc, err := prometheusmetrics.AgentStats(ctx, slogtest.Make(t, nil), registry, db, time.Now().Add(-time.Minute), time.Millisecond)
 	require.NoError(t, err)
-	t.Cleanup(cancel)
+	t.Cleanup(closeFunc)
 
 	// then
 	goldenFile, err := os.ReadFile("testdata/agent-stats.json")
@@ -460,7 +460,7 @@ func TestAgentStats(t *testing.T) {
 	assert.Equal(t, string(goldenFile), string(out))
 }
 
-func prepareWorkspaceAndAgent(t *testing.T, client *codersdk.Client, user codersdk.CreateFirstUserResponse, workspaceNum int) (*agentsdk.Client, codersdk.Workspace) {
+func prepareWorkspaceAndAgent(t *testing.T, client *codersdk.Client, user codersdk.CreateFirstUserResponse, workspaceNum int) *agentsdk.Client {
 	authToken := uuid.NewString()
 
 	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
@@ -477,6 +477,5 @@ func prepareWorkspaceAndAgent(t *testing.T, client *codersdk.Client, user coders
 
 	agentClient := agentsdk.New(client.URL)
 	agentClient.SetSessionToken(authToken)
-
-	return agentClient, workspace
+	return agentClient
 }