chore: improve performance of 'GetLatestWorkspaceBuildsByWorkspaceIDs' (#19452)

Closes https://github.com/coder/internal/issues/716

This prevents a scan over the entire `workspace_build` table by removing
a `join`. This is still imperfect as we are still scanning over the
number of builds for the workspaces in the arguments. Ideally we would
have some index or something precomputed. Then we could skip scanning
over the builds for the correct workspaces that are not the latest.

coderd/database/querier_test.go

@@ -32,6 +32,7 @@ import (
 	"github.com/coder/coder/v2/coderd/provisionerdserver"
 	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/coderd/rbac/policy"
+	"github.com/coder/coder/v2/coderd/util/slice"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/provisionersdk"
 	"github.com/coder/coder/v2/testutil"
@@ -6579,3 +6580,75 @@ func TestWorkspaceBuildDeadlineConstraint(t *testing.T) {
 		}
 	}
 }
+
+// TestGetLatestWorkspaceBuildsByWorkspaceIDs populates the database with
+// workspaces and builds. It then tests that
+// GetLatestWorkspaceBuildsByWorkspaceIDs returns the latest build for some
+// subset of the workspaces.
+func TestGetLatestWorkspaceBuildsByWorkspaceIDs(t *testing.T) {
+	t.Parallel()
+
+	db, _ := dbtestutil.NewDB(t)
+
+	org := dbgen.Organization(t, db, database.Organization{})
+	admin := dbgen.User(t, db, database.User{})
+
+	tv := dbfake.TemplateVersion(t, db).
+		Seed(database.TemplateVersion{
+			OrganizationID: org.ID,
+			CreatedBy:      admin.ID,
+		}).
+		Do()
+
+	users := make([]database.User, 5)
+	wrks := make([][]database.WorkspaceTable, len(users))
+	exp := make(map[uuid.UUID]database.WorkspaceBuild)
+	for i := range users {
+		users[i] = dbgen.User(t, db, database.User{})
+		dbgen.OrganizationMember(t, db, database.OrganizationMember{
+			UserID:         users[i].ID,
+			OrganizationID: org.ID,
+		})
+
+		// Each user gets 2 workspaces.
+		wrks[i] = make([]database.WorkspaceTable, 2)
+		for wi := range wrks[i] {
+			wrks[i][wi] = dbgen.Workspace(t, db, database.WorkspaceTable{
+				TemplateID: tv.Template.ID,
+				OwnerID:    users[i].ID,
+			})
+
+			// Choose a deterministic number of builds per workspace
+			// No more than 5 builds though, that would be excessive.
+			for j := int32(1); int(j) <= (i+wi)%5; j++ {
+				wb := dbfake.WorkspaceBuild(t, db, wrks[i][wi]).
+					Seed(database.WorkspaceBuild{
+						WorkspaceID: wrks[i][wi].ID,
+						BuildNumber: j + 1,
+					}).
+					Do()
+
+				exp[wrks[i][wi].ID] = wb.Build // Save the final workspace build
+			}
+		}
+	}
+
+	// Only take half the users. And only take 1 workspace per user for the test.
+	// The others are just noice. This just queries a subset of workspaces and builds
+	// to make sure the noise doesn't interfere with the results.
+	assertWrks := wrks[:len(users)/2]
+	ctx := testutil.Context(t, testutil.WaitLong)
+	ids := slice.Convert[[]database.WorkspaceTable, uuid.UUID](assertWrks, func(pair []database.WorkspaceTable) uuid.UUID {
+		return pair[0].ID
+	})
+
+	require.Greater(t, len(ids), 0, "expected some workspace ids for test")
+	builds, err := db.GetLatestWorkspaceBuildsByWorkspaceIDs(ctx, ids)
+	require.NoError(t, err)
+	for _, b := range builds {
+		expB, ok := exp[b.WorkspaceID]
+		require.Truef(t, ok, "unexpected workspace build for workspace id %s", b.WorkspaceID)
+		require.Equalf(t, expB.ID, b.ID, "unexpected workspace build id for workspace id %s", b.WorkspaceID)
+		require.Equal(t, expB.BuildNumber, b.BuildNumber, "unexpected build number")
+	}
+}

coderd/database/queries.sql.go

@@ -18983,20 +18983,15 @@ func (q *sqlQuerier) GetLatestWorkspaceBuildByWorkspaceID(ctx context.Context, w
 }
 
 const getLatestWorkspaceBuildsByWorkspaceIDs = `-- name: GetLatestWorkspaceBuildsByWorkspaceIDs :many
-SELECT wb.id, wb.created_at, wb.updated_at, wb.workspace_id, wb.template_version_id, wb.build_number, wb.transition, wb.initiator_id, wb.provisioner_state, wb.job_id, wb.deadline, wb.reason, wb.daily_cost, wb.max_deadline, wb.template_version_preset_id, wb.has_ai_task, wb.ai_task_sidebar_app_id, wb.has_external_agent, wb.initiator_by_avatar_url, wb.initiator_by_username, wb.initiator_by_name
-FROM (
-    SELECT
-        workspace_id, MAX(build_number) as max_build_number
-    FROM
-		workspace_build_with_user AS workspace_builds
-    WHERE
-        workspace_id = ANY($1 :: uuid [ ])
-    GROUP BY
-        workspace_id
-) m
-JOIN
-	 workspace_build_with_user AS wb
-ON m.workspace_id = wb.workspace_id AND m.max_build_number = wb.build_number
+SELECT
+	DISTINCT ON (workspace_id)
+	id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost, max_deadline, template_version_preset_id, has_ai_task, ai_task_sidebar_app_id, has_external_agent, initiator_by_avatar_url, initiator_by_username, initiator_by_name
+FROM
+	workspace_build_with_user AS workspace_builds
+WHERE
+	workspace_id = ANY($1 :: uuid [ ])
+ORDER BY
+	workspace_id, build_number DESC -- latest first
 `
 
 func (q *sqlQuerier) GetLatestWorkspaceBuildsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceBuild, error) {

coderd/database/queries/workspacebuilds.sql

@@ -76,20 +76,16 @@ LIMIT
 	1;
 
 -- name: GetLatestWorkspaceBuildsByWorkspaceIDs :many
-SELECT wb.*
-FROM (
-    SELECT
-        workspace_id, MAX(build_number) as max_build_number
-    FROM
-		workspace_build_with_user AS workspace_builds
-    WHERE
-        workspace_id = ANY(@ids :: uuid [ ])
-    GROUP BY
-        workspace_id
-) m
-JOIN
-	 workspace_build_with_user AS wb
-ON m.workspace_id = wb.workspace_id AND m.max_build_number = wb.build_number;
+SELECT
+	DISTINCT ON (workspace_id)
+	*
+FROM
+	workspace_build_with_user AS workspace_builds
+WHERE
+	workspace_id = ANY(@ids :: uuid [ ])
+ORDER BY
+	workspace_id, build_number DESC -- latest first
+;
 
 -- name: InsertWorkspaceBuild :exec
 INSERT INTO