From b904d0c3816aabce19c2317b30c2d2ec83d6f037 Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Tue, 21 Jan 2025 12:20:41 +0000 Subject: [PATCH] fix: return only the first workspace agent script timing per script (#16203) Fixes https://github.com/coder/coder/issues/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. --- coderd/database/dbgen/dbgen.go | 14 ++++-- coderd/database/dbmem/dbmem.go | 9 ++++ coderd/database/queries.sql.go | 3 +- coderd/database/queries/workspaceagents.sql | 5 +- coderd/workspacebuilds_test.go | 52 ++++++++++++++++++++- coderd/workspaces_test.go | 4 +- 6 files changed, 77 insertions(+), 10 deletions(-) diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 17d3d199639cc..75a958fcc8968 100644 --- 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 index 507f040abbd9b..50e9a1bac055b 100644 --- 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 index ea2b7be288adb..ec75a8ed4f56e 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -11656,7 +11656,7 @@ func (q *sqlQuerier) GetWorkspaceAgentMetadata(ctx context.Context, arg GetWorks const getWorkspaceAgentScriptTimingsByBuildID = `-- name: GetWorkspaceAgentScriptTimingsByBuildID :many SELECT - workspace_agent_script_timings.script_id, workspace_agent_script_timings.started_at, workspace_agent_script_timings.ended_at, workspace_agent_script_timings.exit_code, workspace_agent_script_timings.stage, workspace_agent_script_timings.status, + DISTINCT ON (workspace_agent_script_timings.script_id) workspace_agent_script_timings.script_id, workspace_agent_script_timings.started_at, workspace_agent_script_timings.ended_at, workspace_agent_script_timings.exit_code, workspace_agent_script_timings.stage, workspace_agent_script_timings.status, workspace_agent_scripts.display_name, workspace_agents.id as workspace_agent_id, workspace_agents.name as workspace_agent_name @@ -11666,6 +11666,7 @@ INNER JOIN workspace_agents ON workspace_agents.id = workspace_agent_scripts.wor 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 +ORDER BY workspace_agent_script_timings.script_id, workspace_agent_script_timings.started_at ` type GetWorkspaceAgentScriptTimingsByBuildIDRow struct { diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql index df7c829861cb2..52d8b5275fc97 100644 --- 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; \ No newline at end of file +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 index feb748ad29250..88599a807e262 100644 --- 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 index 6a2856dcbbe76..bc4cdab6b2b65 100644 --- 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)