fix: return only the first workspace agent script timing per script (#16203)

DanielleMaywood · web-flow · commit b904d0c3816a · 2025-01-21T12:20:41.000Z
Fixes #16124 If a workspace agent crashes, it is possible for any startup scripts to be ran again. This PR makes it so that the `GetWorkspaceAgentScriptTimingsByBuildID` query only returns the first timing recorded per-script.
diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -209,9 +209,17 @@ func WorkspaceAgentScript(t testing.TB, db database.Store, orig database.Workspa
 	return scripts[0]
 }
 
-func WorkspaceAgentScriptTimings(t testing.TB, db database.Store, script database.WorkspaceAgentScript, count int) []database.WorkspaceAgentScriptTiming {
-	timings := make([]database.WorkspaceAgentScriptTiming, count)
-	for i := range count {
+func WorkspaceAgentScripts(t testing.TB, db database.Store, count int, orig database.WorkspaceAgentScript) []database.WorkspaceAgentScript {
+	scripts := make([]database.WorkspaceAgentScript, 0, count)
+	for range count {
+		scripts = append(scripts, WorkspaceAgentScript(t, db, orig))
+	}
+	return scripts
+}
+
+func WorkspaceAgentScriptTimings(t testing.TB, db database.Store, scripts []database.WorkspaceAgentScript) []database.WorkspaceAgentScriptTiming {
+	timings := make([]database.WorkspaceAgentScriptTiming, len(scripts))
+	for i, script := range scripts {
 		timings[i] = WorkspaceAgentScriptTiming(t, db, database.WorkspaceAgentScriptTiming{
 			ScriptID: script.ID,
 		})
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -6018,6 +6018,15 @@ func (q *FakeQuerier) GetWorkspaceAgentScriptTimingsByBuildID(ctx context.Contex
 			WorkspaceAgentName: agent.Name,
 		})
 	}
+
+	// We want to only return the first script run for each Script ID.
+	slices.SortFunc(rows, func(a, b database.GetWorkspaceAgentScriptTimingsByBuildIDRow) int {
+		return a.StartedAt.Compare(b.StartedAt)
+	})
+	rows = slices.CompactFunc(rows, func(e1, e2 database.GetWorkspaceAgentScriptTimingsByBuildIDRow) bool {
+		return e1.ScriptID == e2.ScriptID
+	})
+
 	return rows, nil
 }
 
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql
@@ -304,7 +304,7 @@ RETURNING workspace_agent_script_timings.*;
 
 -- name: GetWorkspaceAgentScriptTimingsByBuildID :many
 SELECT
-	workspace_agent_script_timings.*,
+	DISTINCT ON (workspace_agent_script_timings.script_id) workspace_agent_script_timings.*,
 	workspace_agent_scripts.display_name,
 	workspace_agents.id as workspace_agent_id,
 	workspace_agents.name as workspace_agent_name
@@ -313,4 +313,5 @@ INNER JOIN workspace_agent_scripts ON workspace_agent_scripts.id = workspace_age
 INNER JOIN workspace_agents ON workspace_agents.id = workspace_agent_scripts.workspace_agent_id
 INNER JOIN workspace_resources ON workspace_resources.id = workspace_agents.resource_id
 INNER JOIN workspace_builds ON workspace_builds.job_id = workspace_resources.job_id
-WHERE workspace_builds.id = $1;
+WHERE workspace_builds.id = $1
+ORDER BY workspace_agent_script_timings.script_id, workspace_agent_script_timings.started_at;
diff --git a/coderd/workspacebuilds_test.go b/coderd/workspacebuilds_test.go
@@ -14,6 +14,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/propagation"
+	"golang.org/x/exp/slices"
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
@@ -1421,6 +1422,47 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 		}
 	})
 
+	t.Run("MultipleTimingsForSameAgentScript", func(t *testing.T) {
+		t.Parallel()
+
+		// Given: a build with multiple timings for the same script
+		build := makeBuild(t)
+		resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+			JobID: build.JobID,
+		})
+		agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: resource.ID,
+		})
+		script := dbgen.WorkspaceAgentScript(t, db, database.WorkspaceAgentScript{
+			WorkspaceAgentID: agent.ID,
+		})
+		timings := make([]database.WorkspaceAgentScriptTiming, 3)
+		scriptStartedAt := dbtime.Now()
+		for i := range timings {
+			timings[i] = dbgen.WorkspaceAgentScriptTiming(t, db, database.WorkspaceAgentScriptTiming{
+				StartedAt: scriptStartedAt,
+				EndedAt:   scriptStartedAt.Add(1 * time.Minute),
+				ScriptID:  script.ID,
+			})
+
+			// Add an hour to the previous "started at" so we can
+			// reliably differentiate the scripts from each other.
+			scriptStartedAt = scriptStartedAt.Add(1 * time.Hour)
+		}
+
+		// When: fetching timings for the build
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		t.Cleanup(cancel)
+		res, err := client.WorkspaceBuildTimings(ctx, build.ID)
+		require.NoError(t, err)
+
+		// Then: return a response with the first agent script timing
+		require.Len(t, res.AgentScriptTimings, 1)
+
+		require.Equal(t, timings[0].StartedAt.UnixMilli(), res.AgentScriptTimings[0].StartedAt.UnixMilli())
+		require.Equal(t, timings[0].EndedAt.UnixMilli(), res.AgentScriptTimings[0].EndedAt.UnixMilli())
+	})
+
 	t.Run("AgentScriptTimings", func(t *testing.T) {
 		t.Parallel()
 
@@ -1432,10 +1474,10 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 		agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
 			ResourceID: resource.ID,
 		})
-		script := dbgen.WorkspaceAgentScript(t, db, database.WorkspaceAgentScript{
+		scripts := dbgen.WorkspaceAgentScripts(t, db, 5, database.WorkspaceAgentScript{
 			WorkspaceAgentID: agent.ID,
 		})
-		agentScriptTimings := dbgen.WorkspaceAgentScriptTimings(t, db, script, 5)
+		agentScriptTimings := dbgen.WorkspaceAgentScriptTimings(t, db, scripts)
 
 		// When: fetching timings for the build
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
@@ -1445,6 +1487,12 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 
 		// Then: return a response with the expected timings
 		require.Len(t, res.AgentScriptTimings, 5)
+		slices.SortFunc(res.AgentScriptTimings, func(a, b codersdk.AgentScriptTiming) int {
+			return a.StartedAt.Compare(b.StartedAt)
+		})
+		slices.SortFunc(agentScriptTimings, func(a, b database.WorkspaceAgentScriptTiming) int {
+			return a.StartedAt.Compare(b.StartedAt)
+		})
 		for i := range res.AgentScriptTimings {
 			timingRes := res.AgentScriptTimings[i]
 			genTiming := agentScriptTimings[i]
diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go
@@ -3757,10 +3757,10 @@ func TestWorkspaceTimings(t *testing.T) {
 		agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
 			ResourceID: resource.ID,
 		})
-		script := dbgen.WorkspaceAgentScript(t, db, database.WorkspaceAgentScript{
+		scripts := dbgen.WorkspaceAgentScripts(t, db, 3, database.WorkspaceAgentScript{
 			WorkspaceAgentID: agent.ID,
 		})
-		dbgen.WorkspaceAgentScriptTimings(t, db, script, 3)
+		dbgen.WorkspaceAgentScriptTimings(t, db, scripts)
 
 		// When: fetching the timings
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
