fix: use NoLock dbmem subqueries for GetWorkspaceAgentScriptTimingsByBuildID (#15405)

ethanndickson · web-flow · commit f0a4de5e1ad2 · 2024-11-06T23:43:58.000+11:00
Closes coder/internal#202. Addresses the problem mentioned in the issue, and also makes sure subtests don't reuse `testing.T` from the main test.
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -1104,6 +1104,19 @@ func (q *FakeQuerier) getOrganizationByIDNoLock(id uuid.UUID) (database.Organiza
 	return database.Organization{}, sql.ErrNoRows
 }
 
+func (q *FakeQuerier) getWorkspaceAgentScriptsByAgentIDsNoLock(ids []uuid.UUID) ([]database.WorkspaceAgentScript, error) {
+	scripts := make([]database.WorkspaceAgentScript, 0)
+	for _, script := range q.workspaceAgentScripts {
+		for _, id := range ids {
+			if script.WorkspaceAgentID == id {
+				scripts = append(scripts, script)
+				break
+			}
+		}
+	}
+	return scripts, nil
+}
+
 func (*FakeQuerier) AcquireLock(_ context.Context, _ int64) error {
 	return xerrors.New("AcquireLock must only be called within a transaction")
 }
@@ -5854,12 +5867,12 @@ func (q *FakeQuerier) GetWorkspaceAgentScriptTimingsByBuildID(ctx context.Contex
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
 
-	build, err := q.GetWorkspaceBuildByID(ctx, id)
+	build, err := q.getWorkspaceBuildByIDNoLock(ctx, id)
 	if err != nil {
 		return nil, xerrors.Errorf("get build: %w", err)
 	}
 
-	resources, err := q.GetWorkspaceResourcesByJobID(ctx, build.JobID)
+	resources, err := q.getWorkspaceResourcesByJobIDNoLock(ctx, build.JobID)
 	if err != nil {
 		return nil, xerrors.Errorf("get resources: %w", err)
 	}
@@ -5868,7 +5881,7 @@ func (q *FakeQuerier) GetWorkspaceAgentScriptTimingsByBuildID(ctx context.Contex
 		resourceIDs = append(resourceIDs, res.ID)
 	}
 
-	agents, err := q.GetWorkspaceAgentsByResourceIDs(ctx, resourceIDs)
+	agents, err := q.getWorkspaceAgentsByResourceIDsNoLock(ctx, resourceIDs)
 	if err != nil {
 		return nil, xerrors.Errorf("get agents: %w", err)
 	}
@@ -5877,7 +5890,7 @@ func (q *FakeQuerier) GetWorkspaceAgentScriptTimingsByBuildID(ctx context.Contex
 		agentIDs = append(agentIDs, agent.ID)
 	}
 
-	scripts, err := q.GetWorkspaceAgentScriptsByAgentIDs(ctx, agentIDs)
+	scripts, err := q.getWorkspaceAgentScriptsByAgentIDsNoLock(agentIDs)
 	if err != nil {
 		return nil, xerrors.Errorf("get scripts: %w", err)
 	}
@@ -5933,16 +5946,7 @@ func (q *FakeQuerier) GetWorkspaceAgentScriptsByAgentIDs(_ context.Context, ids
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
 
-	scripts := make([]database.WorkspaceAgentScript, 0)
-	for _, script := range q.workspaceAgentScripts {
-		for _, id := range ids {
-			if script.WorkspaceAgentID == id {
-				scripts = append(scripts, script)
-				break
-			}
-		}
-	}
-	return scripts, nil
+	return q.getWorkspaceAgentScriptsByAgentIDsNoLock(ids)
 }
 
 func (q *FakeQuerier) GetWorkspaceAgentStats(_ context.Context, createdAfter time.Time) ([]database.GetWorkspaceAgentStatsRow, error) {
diff --git a/coderd/workspacebuilds_test.go b/coderd/workspacebuilds_test.go
@@ -1219,7 +1219,7 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 
 	// Tests will run in parallel. To avoid conflicts and race conditions on the
 	// build number, each test will have its own workspace and build.
-	makeBuild := func() database.WorkspaceBuild {
+	makeBuild := func(t *testing.T) database.WorkspaceBuild {
 		ws := dbgen.Workspace(t, db, database.WorkspaceTable{
 			OwnerID:        user.ID,
 			OrganizationID: owner.OrganizationID,
@@ -1260,7 +1260,7 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 		t.Parallel()
 
 		// Given: a build with no timings
-		build := makeBuild()
+		build := makeBuild(t)
 
 		// When: fetching timings for the build
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
@@ -1277,7 +1277,7 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 		t.Parallel()
 
 		// Given: a build with provisioner timings
-		build := makeBuild()
+		build := makeBuild(t)
 		provisionerTimings := dbgen.ProvisionerJobTimings(t, db, build, 5)
 
 		// When: fetching timings for the build
@@ -1305,7 +1305,7 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 		t.Parallel()
 
 		// Given: a build with agent script timings
-		build := makeBuild()
+		build := makeBuild(t)
 		resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
 			JobID: build.JobID,
 		})
@@ -1342,7 +1342,7 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 		t.Parallel()
 
 		// Given: a build with no agent scripts
-		build := makeBuild()
+		build := makeBuild(t)
 		resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
 			JobID: build.JobID,
 		})
@@ -1365,7 +1365,7 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 		t.Parallel()
 
 		// Given: a build with no agents
-		build := makeBuild()
+		build := makeBuild(t)
 		dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
 			JobID: build.JobID,
 		})
@@ -1385,7 +1385,7 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 		t.Parallel()
 
 		// Given: a build with an agent
-		build := makeBuild()
+		build := makeBuild(t)
 		resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
 			JobID: build.JobID,
 		})
@@ -1414,7 +1414,7 @@ func TestWorkspaceBuildTimings(t *testing.T) {
 		t.Parallel()
 
 		// Given: a build with multiple agents
-		build := makeBuild()
+		build := makeBuild(t)
 		resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
 			JobID: build.JobID,
 		})
