From dce750db6cfc63592dd53c97b54315d9fc37687f Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Wed, 23 Oct 2024 11:57:07 -0500 Subject: [PATCH 01/19] chore: making unit tests to try out various things --- coderd/database/dbtestutil/tx.go | 65 +++++ enterprise/coderd/coderd.go | 2 +- enterprise/coderd/workspacequota.go | 4 +- enterprise/coderd/workspacequota_test.go | 343 +++++++++++++++++++++++ 4 files changed, 411 insertions(+), 3 deletions(-) create mode 100644 coderd/database/dbtestutil/tx.go diff --git a/coderd/database/dbtestutil/tx.go b/coderd/database/dbtestutil/tx.go new file mode 100644 index 0000000000000..cc8ef2ff0d63a --- /dev/null +++ b/coderd/database/dbtestutil/tx.go @@ -0,0 +1,65 @@ +package dbtestutil + +import ( + "database/sql" + "sync" + "testing" + + "github.com/coder/coder/v2/coderd/database" +) + +type DBTx struct { + database.Store + mu sync.Mutex + err error + errC chan error + finalErr chan error +} + +// StartTx starts a transaction and returns a DBTx object. This allows running +// 2 transactions concurrently in a test more easily. +func StartTx(t *testing.T, db database.Store, opts *sql.TxOptions) *DBTx { + errC := make(chan error) + finalErr := make(chan error) + txC := make(chan database.Store) + + go func() { + t.Helper() + once := sync.Once{} + count := 0 + + err := db.InTx(func(store database.Store) error { + // InTx can be retried + once.Do(func() { + txC <- store + }) + count++ + if count > 1 { + t.Logf("InTx called more than once: %d", count) + } + return <-errC + }, opts) + finalErr <- err + }() + + txStore := <-txC + close(txC) + + return &DBTx{Store: txStore, errC: errC, finalErr: finalErr} +} + +func (tx *DBTx) SetError(err error) { + tx.mu.Lock() + defer tx.mu.Unlock() + tx.err = err +} + +// Done can only be called once. If you call it twice, it will panic. +func (tx *DBTx) Done() error { + tx.mu.Lock() + defer tx.mu.Unlock() + + tx.errC <- tx.err + close(tx.errC) + return <-tx.finalErr +} diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 7e59eb341411f..24dc43f28b576 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -701,7 +701,7 @@ func (api *API) updateEntitlements(ctx context.Context) error { if initial, changed, enabled := featureChanged(codersdk.FeatureTemplateRBAC); shouldUpdate(initial, changed, enabled) { if enabled { - committer := committer{ + committer := Committer{ Log: api.Logger.Named("quota_committer"), Database: api.Database, } diff --git a/enterprise/coderd/workspacequota.go b/enterprise/coderd/workspacequota.go index 7ea42ea24f491..ae9ffe14b6d89 100644 --- a/enterprise/coderd/workspacequota.go +++ b/enterprise/coderd/workspacequota.go @@ -18,12 +18,12 @@ import ( "github.com/coder/coder/v2/provisionerd/proto" ) -type committer struct { +type Committer struct { Log slog.Logger Database database.Store } -func (c *committer) CommitQuota( +func (c *Committer) CommitQuota( ctx context.Context, request *proto.CommitQuotaRequest, ) (*proto.CommitQuotaResponse, error) { jobID, err := uuid.Parse(request.JobId) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index ac4a77eaec8b4..c218cd4257779 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -2,18 +2,26 @@ package coderd_test import ( "context" + "database/sql" "encoding/json" "fmt" "net/http" "sync" "testing" + "time" "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" @@ -295,6 +303,166 @@ func TestWorkspaceQuota(t *testing.T) { }) } +// DB=ci DB_FROM=cikggwjxbths +func TestWorkspaceSerialization(t *testing.T) { + t.Parallel() + + if !dbtestutil.WillUsePostgres() { + panic("We should only run this test with postgres") + } + + db, ps := dbtestutil.NewDB(t) + //c := &coderd.Committer{ + // Log: slogtest.Make(t, nil), + // Database: db, + //} + var _ = ps + + org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + tpl := dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + tplVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + TemplateID: uuid.NullUUID{ + UUID: tpl.ID, + Valid: true, + }, + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + + seed := database.WorkspaceBuild{ + TemplateVersionID: tplVersion.ID, + } + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OrganizationID: org.ID, + OwnerID: user.ID, + TemplateID: tpl.ID, + }) + + workspaceResp := dbfake.WorkspaceBuild(t, db, workspace).Seed(seed).Do() + + workspaceTwo := dbgen.Workspace(t, db, database.WorkspaceTable{ + OrganizationID: org.ID, + OwnerID: user.ID, + TemplateID: tpl.ID, + }) + workspaceTwoResp := dbfake.WorkspaceBuild(t, db, workspaceTwo).Seed(seed).Do() + + // TX mixing tests. **DO NOT** run these in parallel. + // The goal here is to mess around with different ordering of + // transactions and queries. + + // UpdateBuildDeadline bumps a workspace deadline while doing a quota + // commit. + // pq: could not serialize access due to concurrent update + // + // Note: This passes if the interrupt is run before 'GetQuota()' + // Passing orders: + // - BeginTX -> Bump! -> GetQuota -> GetAllowance -> UpdateCost -> EndTx + // - BeginTX -> GetQuota -> GetAllowance -> UpdateCost -> Bump! -> EndTx + t.Run("UpdateBuildDeadline", func(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + bumpDeadline := func() { + err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ + Deadline: dbtime.Now(), + MaxDeadline: dbtime.Now(), + UpdatedAt: dbtime.Now(), + ID: workspaceResp.Build.ID, + }) + require.NoError(t, err) + } + + // Start TX + // Run order + + quota := newCommitter(t, db, workspace, workspaceResp.Build) + quota.GetQuota(ctx, t) // Step 1 + bumpDeadline() // Interrupt + quota.GetAllowance(ctx, t) // Step 2 + quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 + // End commit + require.NoError(t, quota.Done()) + }) + + t.Run("ReadCost", func(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + readCost := func() { + _, err := db.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ + OwnerID: workspace.OwnerID, + OrganizationID: workspace.OrganizationID, + }) + require.NoError(t, err) + } + + // Start TX + // Run order + + quota := newCommitter(t, db, workspace, workspaceResp.Build) + quota.GetQuota(ctx, t) // Step 1 + readCost() // Interrupt + quota.GetAllowance(ctx, t) // Step 2 + quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 + + // End commit + require.NoError(t, quota.Done()) + }) + + t.Run("AutoBuild", func(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + auto := newautobuild(t, db, workspace, workspaceResp.Build) + quota := newCommitter(t, db, workspace, workspaceResp.Build) + + // Run order + auto.DoAllReads(ctx, t) + + quota.GetQuota(ctx, t) // Step 1 + auto.DoAllReads(ctx, t) + + quota.GetAllowance(ctx, t) // Step 2 + auto.DoAllReads(ctx, t) + + quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 + auto.DoAllReads(ctx, t) + auto.DoAllWrites(ctx, t) + + // End commit + require.NoError(t, auto.Done()) + require.NoError(t, quota.Done()) + }) + + t.Run("DoubleCommit", func(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + one := newCommitter(t, db, workspace, workspaceResp.Build) + two := newCommitter(t, db, workspaceTwo, workspaceTwoResp.Build) + + var _, _ = one, two + // Run order + two.GetQuota(ctx, t) + two.GetAllowance(ctx, t) + + one.GetQuota(ctx, t) + one.GetAllowance(ctx, t) + + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + two.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + // End commit + require.NoError(t, one.Done()) + require.NoError(t, two.Done()) + }) +} + func deprecatedQuotaEndpoint(ctx context.Context, client *codersdk.Client, userID string) (codersdk.WorkspaceQuota, error) { res, err := client.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspace-quota/%s", userID), nil) if err != nil { @@ -335,3 +503,178 @@ func applyWithCost(cost int32) []*proto.Response { }, }} } + +type autobuilder struct { + DBTx *dbtestutil.DBTx + w database.WorkspaceTable + b database.WorkspaceBuild + step int + + // some stuff in mem + lastestBuild database.WorkspaceBuild + newJob database.ProvisionerJob +} + +func newautobuild(t *testing.T, db database.Store, workspace database.WorkspaceTable, b database.WorkspaceBuild) *autobuilder { + quotaTX := dbtestutil.StartTx(t, db, &sql.TxOptions{ + Isolation: sql.LevelRepeatableRead, + }) + return &autobuilder{DBTx: quotaTX, w: workspace, b: b} +} + +func (c *autobuilder) DoAllWrites(ctx context.Context, t *testing.T) { + c.InsertProvisionerJob(ctx, t) + c.InsertWorkspaceBuild(ctx, t) +} + +func (c *autobuilder) InsertProvisionerJob(ctx context.Context, t *testing.T) { + now := dbtime.Now() + job, err := c.DBTx.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ + ID: uuid.New(), + CreatedAt: now, + UpdatedAt: now, + OrganizationID: c.w.OrganizationID, + InitiatorID: c.w.OwnerID, + Provisioner: database.ProvisionerTypeTerraform, + StorageMethod: database.ProvisionerStorageMethodFile, + FileID: uuid.New(), + Type: database.ProvisionerJobTypeWorkspaceBuild, + Input: []byte("{}"), + Tags: nil, + TraceMetadata: pqtype.NullRawMessage{}, + }) + require.NoError(t, err) + c.newJob = job +} + +func (c *autobuilder) InsertWorkspaceBuild(ctx context.Context, t *testing.T) { + now := dbtime.Now() + err := c.DBTx.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{ + ID: uuid.New(), + CreatedAt: now, + UpdatedAt: now, + WorkspaceID: c.w.ID, + TemplateVersionID: c.b.TemplateVersionID, + BuildNumber: c.lastestBuild.BuildNumber + 1, + Transition: database.WorkspaceTransitionStart, + InitiatorID: c.w.OwnerID, + JobID: c.newJob.ID, + ProvisionerState: nil, + Deadline: dbtime.Now().Add(time.Hour), + MaxDeadline: dbtime.Now().Add(time.Hour), + Reason: database.BuildReasonAutostart, + }) + require.NoError(t, err) +} + +func (c *autobuilder) DoAllReads(ctx context.Context, t *testing.T) { + c.GetWorkspace(ctx, t) + c.GetUser(ctx, t) + c.GetLatestWorkspaceBuildByWorkspaceID(ctx, t) + c.GetProvisionerJobByID(ctx, t) + c.GetTemplateByID(ctx, t) +} + +func (c *autobuilder) Next(ctx context.Context, t *testing.T) { + list := c.steps() + list[c.step%len(list)](ctx, t) + c.step++ +} + +func (c *autobuilder) steps() []func(context.Context, *testing.T) { + return []func(context.Context, *testing.T){ + noReturn(c.GetWorkspace), + noReturn(c.GetUser), + noReturn(c.GetLatestWorkspaceBuildByWorkspaceID), + noReturn(c.GetProvisionerJobByID), + noReturn(c.GetTemplateByID), + } +} + +func (c *autobuilder) GetWorkspace(ctx context.Context, t *testing.T) database.Workspace { + workspace, err := c.DBTx.GetWorkspaceByID(ctx, c.w.ID) + require.NoError(t, err) + return workspace +} + +func (c *autobuilder) GetUser(ctx context.Context, t *testing.T) database.User { + user, err := c.DBTx.GetUserByID(ctx, c.w.OwnerID) + require.NoError(t, err) + return user +} + +func (c *autobuilder) GetLatestWorkspaceBuildByWorkspaceID(ctx context.Context, t *testing.T) database.WorkspaceBuild { + build, err := c.DBTx.GetLatestWorkspaceBuildByWorkspaceID(ctx, c.w.ID) + require.NoError(t, err) + c.lastestBuild = build + return build +} + +func (c *autobuilder) GetProvisionerJobByID(ctx context.Context, t *testing.T) database.ProvisionerJob { + job, err := c.DBTx.GetProvisionerJobByID(ctx, c.lastestBuild.JobID) + require.NoError(t, err) + return job +} + +func (c *autobuilder) GetTemplateByID(ctx context.Context, t *testing.T) database.Template { + tpl, err := c.DBTx.GetTemplateByID(ctx, c.w.TemplateID) + require.NoError(t, err) + return tpl +} + +func (c *autobuilder) Done() error { + return c.DBTx.Done() +} + +// committer does what the CommitQuota does, but allows +// stepping through the actions in the tx and controlling the +// timing. +type committer struct { + DBTx *dbtestutil.DBTx + w database.WorkspaceTable + b database.WorkspaceBuild +} + +func newCommitter(t *testing.T, db database.Store, workspace database.WorkspaceTable, build database.WorkspaceBuild) *committer { + quotaTX := dbtestutil.StartTx(t, db, &sql.TxOptions{ + Isolation: sql.LevelSerializable, + ReadOnly: false, + }) + return &committer{DBTx: quotaTX, w: workspace, b: build} +} + +func (c *committer) GetQuota(ctx context.Context, t *testing.T) int64 { + consumed, err := c.DBTx.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ + OwnerID: c.w.OwnerID, + OrganizationID: c.w.OrganizationID, + }) + require.NoError(t, err) + return consumed +} + +func (c *committer) GetAllowance(ctx context.Context, t *testing.T) int64 { + allowance, err := c.DBTx.GetQuotaAllowanceForUser(ctx, database.GetQuotaAllowanceForUserParams{ + UserID: c.w.OwnerID, + OrganizationID: c.w.OrganizationID, + }) + require.NoError(t, err) + return allowance +} + +func (c *committer) UpdateWorkspaceBuildCostByID(ctx context.Context, t *testing.T, cost int32) { + err := c.DBTx.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{ + ID: c.b.ID, + DailyCost: cost, + }) + require.NoError(t, err) +} + +func (c *committer) Done() error { + return c.DBTx.Done() +} + +func noReturn[T any](f func(context.Context, *testing.T) T) func(context.Context, *testing.T) { + return func(ctx context.Context, t *testing.T) { + f(ctx, t) + } +} From 4fd97605d6279fd33972e13f45665f1dcc119458 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Wed, 23 Oct 2024 14:16:09 -0500 Subject: [PATCH 02/19] wip --- enterprise/coderd/workspacequota_test.go | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index c218cd4257779..ca0f4d7f545e3 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -458,9 +458,14 @@ func TestWorkspaceSerialization(t *testing.T) { two.UpdateWorkspaceBuildCostByID(ctx, t, 10) // End commit - require.NoError(t, one.Done()) - require.NoError(t, two.Done()) + err := one.Done() + err2 := two.Done() + require.NoError(t, err) + require.NoError(t, err2) }) + + // TODO: Try to fail a non-repeatable read only transaction + // Autobuild, then quota, then autobuild read agin in the same tx } func deprecatedQuotaEndpoint(ctx context.Context, client *codersdk.Client, userID string) (codersdk.WorkspaceQuota, error) { From 62156037f0bb7fd8a8f33f57349ecc43ee588e82 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Thu, 24 Oct 2024 09:13:42 -0500 Subject: [PATCH 03/19] wip --- enterprise/coderd/workspacequota_test.go | 140 ++++++++++++++++++++++- 1 file changed, 135 insertions(+), 5 deletions(-) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index ca0f4d7f545e3..27283199b41a9 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -319,7 +319,31 @@ func TestWorkspaceSerialization(t *testing.T) { var _ = ps org := dbgen.Organization(t, db, database.Organization{}) + + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + group, err := db.InsertAllUsersGroup(ctx, org.ID) + require.NoError(t, err) + + group, err = db.UpdateGroupByID(ctx, database.UpdateGroupByIDParams{ + Name: group.Name, + DisplayName: group.DisplayName, + AvatarURL: group.AvatarURL, + QuotaAllowance: 100, + ID: group.ID, + }) + require.NoError(t, err) + user := dbgen.User(t, db, database.User{}) + + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + CreatedAt: dbtime.Now(), + UpdatedAt: dbtime.Now(), + Roles: []string{}, + }) + tpl := dbgen.Template(t, db, database.Template{ OrganizationID: org.ID, CreatedBy: user.ID, @@ -448,24 +472,130 @@ func TestWorkspaceSerialization(t *testing.T) { var _, _ = one, two // Run order + one.GetQuota(ctx, t) + one.GetAllowance(ctx, t) + + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + two.UpdateWorkspaceBuildCostByID(ctx, t, 10) + two.GetQuota(ctx, t) two.GetAllowance(ctx, t) + // End commit + require.NoError(t, one.Done()) + require.NoError(t, two.Done()) + }) + + t.Run("BumpLastUsedAt", func(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitShort) + ctx = dbauthz.AsSystemRestricted(ctx) + + two := dbtestutil.StartTx(t, db, &sql.TxOptions{ + Isolation: sql.LevelSerializable, + ReadOnly: false, + }) + one := newCommitter(t, db, workspace, workspaceResp.Build) + + //two := newCommitter(t, db, workspaceTwo, workspaceResp.Build) + + // Run order one.GetQuota(ctx, t) one.GetAllowance(ctx, t) + //two.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + //err := two.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{ + // ID: workspaceResp.Build.ID, + // DailyCost: 30, + //}) + //require.NoError(t, err) + + //q, err := two.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ + // OwnerID: user.ID, + // OrganizationID: workspace.OrganizationID, + //}) + //require.NoError(t, err) + //fmt.Println(q) + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - two.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + wg := sync.WaitGroup{} + wg.Add(1) + go func() { + defer wg.Done() + err = two.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ + Deadline: dbtime.Now(), + MaxDeadline: dbtime.Now(), + UpdatedAt: dbtime.Now(), + ID: workspaceResp.Build.ID, + }) + require.NoError(t, err) + }() + time.Sleep(time.Millisecond * 800) + + //err = db.UpdateWorkspaceLastUsedAt(ctx, database.UpdateWorkspaceLastUsedAtParams{ + // ID: workspace.ID, + // LastUsedAt: dbtime.Now(), + //}) + //require.NoError(t, err) + // + //err = db.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{ + // ID: workspaceResp.Build.ID, + // DailyCost: 20, + //}) + //require.NoError(t, err) + + //two.GetQuota(ctx, t) + //two.GetAllowance(ctx, t) // End commit - err := one.Done() - err2 := two.Done() - require.NoError(t, err) - require.NoError(t, err2) + require.NoError(t, one.Done()) + wg.Wait() + require.NoError(t, two.Done()) }) // TODO: Try to fail a non-repeatable read only transaction + t.Run("ReadStale", func(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + three := newCommitter(t, db, workspace, workspaceTwoResp.Build) + two := newCommitter(t, db, workspaceTwo, workspaceResp.Build) + one := newCommitter(t, db, workspace, workspaceResp.Build) + three.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + // Run order + + fmt.Println("1", one.GetQuota(ctx, t)) + one.GetAllowance(ctx, t) + + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + fmt.Println("1a", one.GetQuota(ctx, t)) + + fmt.Println("2a", two.GetQuota(ctx, t)) + two.GetAllowance(ctx, t) + + // End commit + require.NoError(t, one.Done()) + + fmt.Println("2a", two.GetQuota(ctx, t)) + two.GetAllowance(ctx, t) + + require.NoError(t, two.Done()) + require.NoError(t, three.Done()) + + //allow, err = db.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ + // OwnerID: user.ID, + // OrganizationID: org.ID, + //}) + //require.NoError(t, err) + //fmt.Println(allow) + }) + // Autobuild, then quota, then autobuild read agin in the same tx + // https://www.richardstrnad.ch/posts/go-sql-how-to-get-detailed-error/ + // https://blog.danslimmon.com/2024/01/10/why-transaction-order-matters-even-if-youre-only-reading/ } func deprecatedQuotaEndpoint(ctx context.Context, client *codersdk.Client, userID string) (codersdk.WorkspaceQuota, error) { From 9ccd0b2b85decf8f68207390fc0250bf0da54d44 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Mon, 28 Oct 2024 17:16:02 -0400 Subject: [PATCH 04/19] wip --- enterprise/coderd/workspacequota_test.go | 283 ++++++++++++++++++----- 1 file changed, 228 insertions(+), 55 deletions(-) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 27283199b41a9..3eb36b498f3e1 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -61,6 +61,10 @@ func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, org func TestWorkspaceQuota(t *testing.T) { t.Parallel() + if !dbtestutil.WillUsePostgres() { + t.Fatal("We should only run this test with postgres") + } + // This first test verifies the behavior of creating and deleting workspaces. // It also tests multi-group quota stacking and the everyone group. t.Run("CreateDelete", func(t *testing.T) { @@ -381,13 +385,71 @@ func TestWorkspaceSerialization(t *testing.T) { // UpdateBuildDeadline bumps a workspace deadline while doing a quota // commit. - // pq: could not serialize access due to concurrent update // // Note: This passes if the interrupt is run before 'GetQuota()' // Passing orders: // - BeginTX -> Bump! -> GetQuota -> GetAllowance -> UpdateCost -> EndTx // - BeginTX -> GetQuota -> GetAllowance -> UpdateCost -> Bump! -> EndTx t.Run("UpdateBuildDeadline", func(t *testing.T) { + // +------------------------------+------------------+ + // | Begin Tx | | + // +------------------------------+------------------+ + // | GetQuota(user) | | + // +------------------------------+------------------+ + // | | BumpDeadline(w1) | + // +------------------------------+------------------+ + // | GetAllowance(user) | | + // +------------------------------+------------------+ + // | UpdateWorkspaceBuildCost(w1) | | + // +------------------------------+------------------+ + // | CommitTx() | | + // +------------------------------+------------------+ + // pq: could not serialize access due to concurrent update + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + bumpDeadline := func() { + err = db.InTx(func(db database.Store) error { + err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ + Deadline: dbtime.Now(), + MaxDeadline: dbtime.Now(), + UpdatedAt: dbtime.Now(), + ID: workspaceResp.Build.ID, + }) + return err + }, &sql.TxOptions{ + Isolation: sql.LevelSerializable, + }) + assert.NoError(t, err) + + } + + // Start TX + // Run order + + quota := newCommitter(t, db, workspace, workspaceResp.Build) + quota.GetQuota(ctx, t) // Step 1 + bumpDeadline() // Interrupt + quota.GetAllowance(ctx, t) // Step 2 + quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 + // End commit + require.NoError(t, quota.Done()) + }) + + t.Run("UpdateOtherBuildDeadline", func(t *testing.T) { + // +------------------------------+------------------+ + // | Begin Tx | | + // +------------------------------+------------------+ + // | GetQuota(user) | | + // +------------------------------+------------------+ + // | | BumpDeadline(w2) | + // +------------------------------+------------------+ + // | GetAllowance(user) | | + // +------------------------------+------------------+ + // | UpdateWorkspaceBuildCost(w1) | | + // +------------------------------+------------------+ + // | CommitTx() | | + // +------------------------------+------------------+ ctx := testutil.Context(t, testutil.WaitLong) ctx = dbauthz.AsSystemRestricted(ctx) @@ -396,7 +458,7 @@ func TestWorkspaceSerialization(t *testing.T) { Deadline: dbtime.Now(), MaxDeadline: dbtime.Now(), UpdatedAt: dbtime.Now(), - ID: workspaceResp.Build.ID, + ID: workspaceTwoResp.Build.ID, }) require.NoError(t, err) } @@ -464,6 +526,30 @@ func TestWorkspaceSerialization(t *testing.T) { }) t.Run("DoubleCommit", func(t *testing.T) { + // +---------------------+---------------------+ + // | W1 Quota Tx | W2 Quota Tx | + // +---------------------+---------------------+ + // | Begin Tx | | + // +---------------------+---------------------+ + // | | Begin Tx | + // +---------------------+---------------------+ + // | GetQuota(w1) | | + // +---------------------+---------------------+ + // | GetAllowance(w1) | | + // +---------------------+---------------------+ + // | UpdateBuildCost(w1) | | + // +---------------------+---------------------+ + // | | UpdateBuildCost(w2) | + // +---------------------+---------------------+ + // | | GetQuota(w2) | + // +---------------------+---------------------+ + // | | GetAllowance(w2) | + // +---------------------+---------------------+ + // | CommitTx() | | + // +---------------------+---------------------+ + // | | CommitTx() | + // +---------------------+---------------------+ + // pq: could not serialize access due to read/write dependencies among transactions ctx := testutil.Context(t, testutil.WaitLong) ctx = dbauthz.AsSystemRestricted(ctx) @@ -476,10 +562,10 @@ func TestWorkspaceSerialization(t *testing.T) { one.GetAllowance(ctx, t) one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - two.UpdateWorkspaceBuildCostByID(ctx, t, 10) two.GetQuota(ctx, t) two.GetAllowance(ctx, t) + two.UpdateWorkspaceBuildCostByID(ctx, t, 10) // End commit require.NoError(t, one.Done()) @@ -487,71 +573,144 @@ func TestWorkspaceSerialization(t *testing.T) { }) t.Run("BumpLastUsedAt", func(t *testing.T) { + // +---------------------+----------------------------------+ + // | W1 Quota Tx | | + // +---------------------+----------------------------------+ + // | Begin Tx | | + // +---------------------+----------------------------------+ + // | GetQuota(w1) | | + // +---------------------+----------------------------------+ + // | GetAllowance(w1) | | + // +---------------------+----------------------------------+ + // | | UpdateWorkspaceBuildDeadline(w1) | + // +---------------------+----------------------------------+ + // | UpdateBuildCost(w1) | | + // +---------------------+----------------------------------+ + // | CommitTx() | | + // +---------------------+----------------------------------+ + // pq: could not serialize access due to concurrent update ctx := testutil.Context(t, testutil.WaitShort) ctx = dbauthz.AsSystemRestricted(ctx) - two := dbtestutil.StartTx(t, db, &sql.TxOptions{ - Isolation: sql.LevelSerializable, - ReadOnly: false, - }) one := newCommitter(t, db, workspace, workspaceResp.Build) - //two := newCommitter(t, db, workspaceTwo, workspaceResp.Build) - // Run order one.GetQuota(ctx, t) one.GetAllowance(ctx, t) - //two.UpdateWorkspaceBuildCostByID(ctx, t, 10) + err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ + Deadline: dbtime.Now(), + MaxDeadline: dbtime.Now(), + UpdatedAt: dbtime.Now(), + ID: workspaceResp.Build.ID, + }) + assert.NoError(t, err) - //err := two.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{ - // ID: workspaceResp.Build.ID, - // DailyCost: 30, - //}) - //require.NoError(t, err) + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - //q, err := two.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ - // OwnerID: user.ID, - // OrganizationID: workspace.OrganizationID, - //}) - //require.NoError(t, err) - //fmt.Println(q) + // End commit + assert.NoError(t, one.Done()) + }) + + t.Run("ActivityBump", func(t *testing.T) { + // +---------------------+----------------------------------+ + // | W1 Quota Tx | | + // +---------------------+----------------------------------+ + // | Begin Tx | | + // +---------------------+----------------------------------+ + // | GetQuota(w1) | | + // +---------------------+----------------------------------+ + // | GetAllowance(w1) | | + // +---------------------+----------------------------------+ + // | | ActivityBump(w1) | + // +---------------------+----------------------------------+ + // | UpdateBuildCost(w1) | | + // +---------------------+----------------------------------+ + // | CommitTx() | | + // +---------------------+----------------------------------+ + // pq: could not serialize access due to concurrent update + ctx := testutil.Context(t, testutil.WaitShort) + ctx = dbauthz.AsSystemRestricted(ctx) + + one := newCommitter(t, db, workspace, workspaceResp.Build) + + // Run order + one.GetQuota(ctx, t) + one.GetAllowance(ctx, t) + + err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ + NextAutostart: time.Now(), + WorkspaceID: workspaceResp.Workspace.ID, + }) + + assert.NoError(t, err) one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - wg := sync.WaitGroup{} - wg.Add(1) - go func() { - defer wg.Done() - err = two.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ - Deadline: dbtime.Now(), - MaxDeadline: dbtime.Now(), - UpdatedAt: dbtime.Now(), - ID: workspaceResp.Build.ID, - }) - require.NoError(t, err) - }() - time.Sleep(time.Millisecond * 800) + // End commit + assert.NoError(t, one.Done()) + }) - //err = db.UpdateWorkspaceLastUsedAt(ctx, database.UpdateWorkspaceLastUsedAtParams{ - // ID: workspace.ID, - // LastUsedAt: dbtime.Now(), - //}) - //require.NoError(t, err) - // - //err = db.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{ - // ID: workspaceResp.Build.ID, - // DailyCost: 20, - //}) - //require.NoError(t, err) + t.Run("UserMod", func(t *testing.T) { + // +---------------------+----------------------------------+ + // | W1 Quota Tx | | + // +---------------------+----------------------------------+ + // | Begin Tx | | + // +---------------------+----------------------------------+ + // | GetQuota(w1) | | + // +---------------------+----------------------------------+ + // | GetAllowance(w1) | | + // +---------------------+----------------------------------+ + // | | UpdateWorkspaceBuildDeadline(w1) | + // +---------------------+----------------------------------+ + // | UpdateBuildCost(w1) | | + // +---------------------+----------------------------------+ + // | CommitTx() | | + // +---------------------+----------------------------------+ + // pq: could not serialize access due to concurrent update + ctx := testutil.Context(t, testutil.WaitShort) + ctx = dbauthz.AsSystemRestricted(ctx) + + db.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{ + OrganizationID: org.ID, + UserID: user.ID, + }) + + one := newCommitter(t, db, workspace, workspaceResp.Build) + + // Run order + + err = db.RemoveUserFromAllGroups(ctx, user.ID) + assert.NoError(t, err) + + err = db.DeleteOrganizationMember(ctx, database.DeleteOrganizationMemberParams{ + OrganizationID: org.ID, + UserID: user.ID, + }) + assert.NoError(t, err) + + err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ + NextAutostart: time.Now(), + WorkspaceID: workspaceResp.Workspace.ID, + }) + + one.GetQuota(ctx, t) + one.GetAllowance(ctx, t) - //two.GetQuota(ctx, t) - //two.GetAllowance(ctx, t) + err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ + NextAutostart: time.Now(), + WorkspaceID: workspaceResp.Workspace.ID, + }) + + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ + NextAutostart: time.Now(), + WorkspaceID: workspaceResp.Workspace.ID, + }) // End commit - require.NoError(t, one.Done()) - wg.Wait() - require.NoError(t, two.Done()) + assert.NoError(t, one.Done()) }) // TODO: Try to fail a non-repeatable read only transaction @@ -577,13 +736,13 @@ func TestWorkspaceSerialization(t *testing.T) { two.GetAllowance(ctx, t) // End commit - require.NoError(t, one.Done()) + assert.NoError(t, one.Done()) fmt.Println("2a", two.GetQuota(ctx, t)) two.GetAllowance(ctx, t) - require.NoError(t, two.Done()) - require.NoError(t, three.Done()) + assert.NoError(t, two.Done()) + assert.NoError(t, three.Done()) //allow, err = db.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ // OwnerID: user.ID, @@ -778,7 +937,12 @@ func newCommitter(t *testing.T, db database.Store, workspace database.WorkspaceT return &committer{DBTx: quotaTX, w: workspace, b: build} } +// GetQuota touches: +// - workspace_builds +// - workspaces func (c *committer) GetQuota(ctx context.Context, t *testing.T) int64 { + t.Helper() + consumed, err := c.DBTx.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ OwnerID: c.w.OwnerID, OrganizationID: c.w.OrganizationID, @@ -787,7 +951,14 @@ func (c *committer) GetQuota(ctx context.Context, t *testing.T) int64 { return consumed } +// GetAllowance touches: +// - group_members_expanded +// - users +// - groups +// - org_members func (c *committer) GetAllowance(ctx context.Context, t *testing.T) int64 { + t.Helper() + allowance, err := c.DBTx.GetQuotaAllowanceForUser(ctx, database.GetQuotaAllowanceForUserParams{ UserID: c.w.OwnerID, OrganizationID: c.w.OrganizationID, @@ -796,12 +967,14 @@ func (c *committer) GetAllowance(ctx context.Context, t *testing.T) int64 { return allowance } -func (c *committer) UpdateWorkspaceBuildCostByID(ctx context.Context, t *testing.T, cost int32) { +func (c *committer) UpdateWorkspaceBuildCostByID(ctx context.Context, t *testing.T, cost int32) bool { + t.Helper() + err := c.DBTx.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{ ID: c.b.ID, DailyCost: cost, }) - require.NoError(t, err) + return assert.NoError(t, err) } func (c *committer) Done() error { From 784a98fe3f56ff6c35774923ef19855b69e427c4 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Mon, 28 Oct 2024 18:44:49 -0400 Subject: [PATCH 05/19] refactor tests --- coderd/database/db.go | 2 +- coderd/database/dbfake/builder.go | 121 +++ coderd/database/dbgen/dbgen.go | 2 + coderd/database/dbtestutil/tx.go | 3 +- enterprise/coderd/workspacequota_test.go | 913 ++++++++++++++--------- 5 files changed, 693 insertions(+), 348 deletions(-) create mode 100644 coderd/database/dbfake/builder.go diff --git a/coderd/database/db.go b/coderd/database/db.go index ae2c31a566cb3..e2eeb10d6fb21 100644 --- a/coderd/database/db.go +++ b/coderd/database/db.go @@ -144,7 +144,7 @@ func (q *sqlQuerier) InTx(function func(Store) error, txOpts *TxOptions) error { // We do not want to duplicate those retries. if !inTx && sqlOpts.Isolation == sql.LevelSerializable { // This is an arbitrarily chosen number. - const retryAmount = 3 + const retryAmount = 1 var err error attempts := 0 for attempts = 0; attempts < retryAmount; attempts++ { diff --git a/coderd/database/dbfake/builder.go b/coderd/database/dbfake/builder.go new file mode 100644 index 0000000000000..c4731c6a4693b --- /dev/null +++ b/coderd/database/dbfake/builder.go @@ -0,0 +1,121 @@ +package dbfake + +import ( + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/testutil" +) + +type OrganizationBuilder struct { + t *testing.T + db database.Store + seed database.Organization + allUsersAllowance int32 + members []uuid.UUID + groups map[database.Group][]uuid.UUID +} + +func Organization(t *testing.T, db database.Store) OrganizationBuilder { + return OrganizationBuilder{ + t: t, + db: db, + members: []uuid.UUID{}, + groups: make(map[database.Group][]uuid.UUID), + } +} + +type OrganizationResponse struct { + Org database.Organization + AllUsersGroup database.Group + Members []database.OrganizationMember + Groups []database.Group +} + +func (b OrganizationBuilder) EveryoneAllowance(allowance int) OrganizationBuilder { + b.allUsersAllowance = int32(allowance) + return b +} + +func (b OrganizationBuilder) Seed(seed database.Organization) OrganizationBuilder { + b.seed = seed + return b +} + +func (b OrganizationBuilder) Members(users ...database.User) OrganizationBuilder { + for _, u := range users { + b.members = append(b.members, u.ID) + } + return b +} + +func (b OrganizationBuilder) Group(seed database.Group, members ...database.User) OrganizationBuilder { + b.groups[seed] = []uuid.UUID{} + for _, u := range members { + b.groups[seed] = append(b.groups[seed], u.ID) + } + return b +} + +func (b OrganizationBuilder) Do() OrganizationResponse { + org := dbgen.Organization(b.t, b.db, b.seed) + + ctx := testutil.Context(b.t, testutil.WaitShort) + ctx = dbauthz.AsSystemRestricted(ctx) + everyone, err := b.db.InsertAllUsersGroup(ctx, org.ID) + require.NoError(b.t, err) + + if b.allUsersAllowance > 0 { + everyone, err = b.db.UpdateGroupByID(ctx, database.UpdateGroupByIDParams{ + Name: everyone.Name, + DisplayName: everyone.DisplayName, + AvatarURL: everyone.AvatarURL, + QuotaAllowance: b.allUsersAllowance, + ID: everyone.ID, + }) + require.NoError(b.t, err) + } + + members := make([]database.OrganizationMember, 0) + if len(b.members) > 0 { + for _, u := range b.members { + newMem := dbgen.OrganizationMember(b.t, b.db, database.OrganizationMember{ + UserID: u, + OrganizationID: org.ID, + CreatedAt: dbtime.Now(), + UpdatedAt: dbtime.Now(), + Roles: nil, + }) + members = append(members, newMem) + } + } + + groups := make([]database.Group, 0) + if len(b.groups) > 0 { + for g, users := range b.groups { + g.OrganizationID = org.ID + group := dbgen.Group(b.t, b.db, g) + groups = append(groups, group) + + for _, u := range users { + dbgen.GroupMember(b.t, b.db, database.GroupMemberTable{ + UserID: u, + GroupID: group.ID, + }) + } + } + } + + return OrganizationResponse{ + Org: org, + AllUsersGroup: everyone, + Members: members, + Groups: groups, + } +} diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 69419b98c79b1..d369d8a023ba9 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -407,6 +407,8 @@ func OrganizationMember(t testing.TB, db database.Store, orig database.Organizat } func Group(t testing.TB, db database.Store, orig database.Group) database.Group { + t.Helper() + name := takeFirst(orig.Name, testutil.GetRandomName(t)) group, err := db.InsertGroup(genCtx, database.InsertGroupParams{ ID: takeFirst(orig.ID, uuid.New()), diff --git a/coderd/database/dbtestutil/tx.go b/coderd/database/dbtestutil/tx.go index cc8ef2ff0d63a..5466657d8f9a0 100644 --- a/coderd/database/dbtestutil/tx.go +++ b/coderd/database/dbtestutil/tx.go @@ -1,7 +1,6 @@ package dbtestutil import ( - "database/sql" "sync" "testing" @@ -18,7 +17,7 @@ type DBTx struct { // StartTx starts a transaction and returns a DBTx object. This allows running // 2 transactions concurrently in a test more easily. -func StartTx(t *testing.T, db database.Store, opts *sql.TxOptions) *DBTx { +func StartTx(t *testing.T, db database.Store, opts *database.TxOptions) *DBTx { errC := make(chan error) finalErr := make(chan error) txC := make(chan database.Store) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 3eb36b498f3e1..c469cf5cb4315 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -316,75 +316,41 @@ func TestWorkspaceSerialization(t *testing.T) { } db, ps := dbtestutil.NewDB(t) - //c := &coderd.Committer{ - // Log: slogtest.Make(t, nil), - // Database: db, - //} var _ = ps - org := dbgen.Organization(t, db, database.Organization{}) - - ctx := testutil.Context(t, testutil.WaitLong) - ctx = dbauthz.AsSystemRestricted(ctx) - group, err := db.InsertAllUsersGroup(ctx, org.ID) - require.NoError(t, err) - - group, err = db.UpdateGroupByID(ctx, database.UpdateGroupByIDParams{ - Name: group.Name, - DisplayName: group.DisplayName, - AvatarURL: group.AvatarURL, - QuotaAllowance: 100, - ID: group.ID, - }) - require.NoError(t, err) - user := dbgen.User(t, db, database.User{}) - - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user.ID, - OrganizationID: org.ID, - CreatedAt: dbtime.Now(), - UpdatedAt: dbtime.Now(), - Roles: []string{}, - }) - - tpl := dbgen.Template(t, db, database.Template{ - OrganizationID: org.ID, - CreatedBy: user.ID, - }) - tplVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{ - TemplateID: uuid.NullUUID{ - UUID: tpl.ID, - Valid: true, - }, - OrganizationID: org.ID, - CreatedBy: user.ID, - }) - - seed := database.WorkspaceBuild{ - TemplateVersionID: tplVersion.ID, - } - workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ - OrganizationID: org.ID, - OwnerID: user.ID, - TemplateID: tpl.ID, - }) - - workspaceResp := dbfake.WorkspaceBuild(t, db, workspace).Seed(seed).Do() - - workspaceTwo := dbgen.Workspace(t, db, database.WorkspaceTable{ - OrganizationID: org.ID, - OwnerID: user.ID, - TemplateID: tpl.ID, - }) - workspaceTwoResp := dbfake.WorkspaceBuild(t, db, workspaceTwo).Seed(seed).Do() + otherUser := dbgen.User(t, db, database.User{}) + + org := dbfake.Organization(t, db). + EveryoneAllowance(20). + Members(user, otherUser). + Group(database.Group{ + QuotaAllowance: 10, + }, user, otherUser). + Group(database.Group{ + QuotaAllowance: 10, + }, user). + Do() + + otherOrg := dbfake.Organization(t, db). + EveryoneAllowance(20). + Members(user, otherUser). + Group(database.Group{ + QuotaAllowance: 10, + }, user, otherUser). + Group(database.Group{ + QuotaAllowance: 10, + }, user). + Do() + + var _ = otherOrg // TX mixing tests. **DO NOT** run these in parallel. // The goal here is to mess around with different ordering of // transactions and queries. // UpdateBuildDeadline bumps a workspace deadline while doing a quota - // commit. + // commit to the same workspace build. // // Note: This passes if the interrupt is run before 'GetQuota()' // Passing orders: @@ -408,16 +374,21 @@ func TestWorkspaceSerialization(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) ctx = dbauthz.AsSystemRestricted(ctx) + myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: org.Org.ID, + OwnerID: user.ID, + }).Do() + bumpDeadline := func() { - err = db.InTx(func(db database.Store) error { + err := db.InTx(func(db database.Store) error { err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ Deadline: dbtime.Now(), MaxDeadline: dbtime.Now(), UpdatedAt: dbtime.Now(), - ID: workspaceResp.Build.ID, + ID: myWorkspace.Build.ID, }) return err - }, &sql.TxOptions{ + }, &database.TxOptions{ Isolation: sql.LevelSerializable, }) assert.NoError(t, err) @@ -427,46 +398,7 @@ func TestWorkspaceSerialization(t *testing.T) { // Start TX // Run order - quota := newCommitter(t, db, workspace, workspaceResp.Build) - quota.GetQuota(ctx, t) // Step 1 - bumpDeadline() // Interrupt - quota.GetAllowance(ctx, t) // Step 2 - quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 - // End commit - require.NoError(t, quota.Done()) - }) - - t.Run("UpdateOtherBuildDeadline", func(t *testing.T) { - // +------------------------------+------------------+ - // | Begin Tx | | - // +------------------------------+------------------+ - // | GetQuota(user) | | - // +------------------------------+------------------+ - // | | BumpDeadline(w2) | - // +------------------------------+------------------+ - // | GetAllowance(user) | | - // +------------------------------+------------------+ - // | UpdateWorkspaceBuildCost(w1) | | - // +------------------------------+------------------+ - // | CommitTx() | | - // +------------------------------+------------------+ - ctx := testutil.Context(t, testutil.WaitLong) - ctx = dbauthz.AsSystemRestricted(ctx) - - bumpDeadline := func() { - err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ - Deadline: dbtime.Now(), - MaxDeadline: dbtime.Now(), - UpdatedAt: dbtime.Now(), - ID: workspaceTwoResp.Build.ID, - }) - require.NoError(t, err) - } - - // Start TX - // Run order - - quota := newCommitter(t, db, workspace, workspaceResp.Build) + quota := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) quota.GetQuota(ctx, t) // Step 1 bumpDeadline() // Interrupt quota.GetAllowance(ctx, t) // Step 2 @@ -475,143 +407,6 @@ func TestWorkspaceSerialization(t *testing.T) { require.NoError(t, quota.Done()) }) - t.Run("ReadCost", func(t *testing.T) { - ctx := testutil.Context(t, testutil.WaitLong) - ctx = dbauthz.AsSystemRestricted(ctx) - - readCost := func() { - _, err := db.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ - OwnerID: workspace.OwnerID, - OrganizationID: workspace.OrganizationID, - }) - require.NoError(t, err) - } - - // Start TX - // Run order - - quota := newCommitter(t, db, workspace, workspaceResp.Build) - quota.GetQuota(ctx, t) // Step 1 - readCost() // Interrupt - quota.GetAllowance(ctx, t) // Step 2 - quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 - - // End commit - require.NoError(t, quota.Done()) - }) - - t.Run("AutoBuild", func(t *testing.T) { - ctx := testutil.Context(t, testutil.WaitLong) - ctx = dbauthz.AsSystemRestricted(ctx) - - auto := newautobuild(t, db, workspace, workspaceResp.Build) - quota := newCommitter(t, db, workspace, workspaceResp.Build) - - // Run order - auto.DoAllReads(ctx, t) - - quota.GetQuota(ctx, t) // Step 1 - auto.DoAllReads(ctx, t) - - quota.GetAllowance(ctx, t) // Step 2 - auto.DoAllReads(ctx, t) - - quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 - auto.DoAllReads(ctx, t) - auto.DoAllWrites(ctx, t) - - // End commit - require.NoError(t, auto.Done()) - require.NoError(t, quota.Done()) - }) - - t.Run("DoubleCommit", func(t *testing.T) { - // +---------------------+---------------------+ - // | W1 Quota Tx | W2 Quota Tx | - // +---------------------+---------------------+ - // | Begin Tx | | - // +---------------------+---------------------+ - // | | Begin Tx | - // +---------------------+---------------------+ - // | GetQuota(w1) | | - // +---------------------+---------------------+ - // | GetAllowance(w1) | | - // +---------------------+---------------------+ - // | UpdateBuildCost(w1) | | - // +---------------------+---------------------+ - // | | UpdateBuildCost(w2) | - // +---------------------+---------------------+ - // | | GetQuota(w2) | - // +---------------------+---------------------+ - // | | GetAllowance(w2) | - // +---------------------+---------------------+ - // | CommitTx() | | - // +---------------------+---------------------+ - // | | CommitTx() | - // +---------------------+---------------------+ - // pq: could not serialize access due to read/write dependencies among transactions - ctx := testutil.Context(t, testutil.WaitLong) - ctx = dbauthz.AsSystemRestricted(ctx) - - one := newCommitter(t, db, workspace, workspaceResp.Build) - two := newCommitter(t, db, workspaceTwo, workspaceTwoResp.Build) - - var _, _ = one, two - // Run order - one.GetQuota(ctx, t) - one.GetAllowance(ctx, t) - - one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - - two.GetQuota(ctx, t) - two.GetAllowance(ctx, t) - two.UpdateWorkspaceBuildCostByID(ctx, t, 10) - - // End commit - require.NoError(t, one.Done()) - require.NoError(t, two.Done()) - }) - - t.Run("BumpLastUsedAt", func(t *testing.T) { - // +---------------------+----------------------------------+ - // | W1 Quota Tx | | - // +---------------------+----------------------------------+ - // | Begin Tx | | - // +---------------------+----------------------------------+ - // | GetQuota(w1) | | - // +---------------------+----------------------------------+ - // | GetAllowance(w1) | | - // +---------------------+----------------------------------+ - // | | UpdateWorkspaceBuildDeadline(w1) | - // +---------------------+----------------------------------+ - // | UpdateBuildCost(w1) | | - // +---------------------+----------------------------------+ - // | CommitTx() | | - // +---------------------+----------------------------------+ - // pq: could not serialize access due to concurrent update - ctx := testutil.Context(t, testutil.WaitShort) - ctx = dbauthz.AsSystemRestricted(ctx) - - one := newCommitter(t, db, workspace, workspaceResp.Build) - - // Run order - one.GetQuota(ctx, t) - one.GetAllowance(ctx, t) - - err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ - Deadline: dbtime.Now(), - MaxDeadline: dbtime.Now(), - UpdatedAt: dbtime.Now(), - ID: workspaceResp.Build.ID, - }) - assert.NoError(t, err) - - one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - - // End commit - assert.NoError(t, one.Done()) - }) - t.Run("ActivityBump", func(t *testing.T) { // +---------------------+----------------------------------+ // | W1 Quota Tx | | @@ -632,125 +427,552 @@ func TestWorkspaceSerialization(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) ctx = dbauthz.AsSystemRestricted(ctx) - one := newCommitter(t, db, workspace, workspaceResp.Build) + myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: org.Org.ID, + OwnerID: user.ID, + }).Do() + + one := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) // Run order one.GetQuota(ctx, t) one.GetAllowance(ctx, t) - err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ + err := db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ NextAutostart: time.Now(), - WorkspaceID: workspaceResp.Workspace.ID, - }) - - assert.NoError(t, err) - - one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - - // End commit - assert.NoError(t, one.Done()) - }) - - t.Run("UserMod", func(t *testing.T) { - // +---------------------+----------------------------------+ - // | W1 Quota Tx | | - // +---------------------+----------------------------------+ - // | Begin Tx | | - // +---------------------+----------------------------------+ - // | GetQuota(w1) | | - // +---------------------+----------------------------------+ - // | GetAllowance(w1) | | - // +---------------------+----------------------------------+ - // | | UpdateWorkspaceBuildDeadline(w1) | - // +---------------------+----------------------------------+ - // | UpdateBuildCost(w1) | | - // +---------------------+----------------------------------+ - // | CommitTx() | | - // +---------------------+----------------------------------+ - // pq: could not serialize access due to concurrent update - ctx := testutil.Context(t, testutil.WaitShort) - ctx = dbauthz.AsSystemRestricted(ctx) - - db.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{ - OrganizationID: org.ID, - UserID: user.ID, + WorkspaceID: myWorkspace.Workspace.ID, }) - one := newCommitter(t, db, workspace, workspaceResp.Build) - - // Run order - - err = db.RemoveUserFromAllGroups(ctx, user.ID) assert.NoError(t, err) - err = db.DeleteOrganizationMember(ctx, database.DeleteOrganizationMemberParams{ - OrganizationID: org.ID, - UserID: user.ID, - }) - assert.NoError(t, err) - - err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ - NextAutostart: time.Now(), - WorkspaceID: workspaceResp.Workspace.ID, - }) - - one.GetQuota(ctx, t) - one.GetAllowance(ctx, t) - - err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ - NextAutostart: time.Now(), - WorkspaceID: workspaceResp.Workspace.ID, - }) - one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ - NextAutostart: time.Now(), - WorkspaceID: workspaceResp.Workspace.ID, - }) - // End commit assert.NoError(t, one.Done()) }) - // TODO: Try to fail a non-repeatable read only transaction - t.Run("ReadStale", func(t *testing.T) { - ctx := testutil.Context(t, testutil.WaitLong) - ctx = dbauthz.AsSystemRestricted(ctx) - - three := newCommitter(t, db, workspace, workspaceTwoResp.Build) - two := newCommitter(t, db, workspaceTwo, workspaceResp.Build) - one := newCommitter(t, db, workspace, workspaceResp.Build) - three.UpdateWorkspaceBuildCostByID(ctx, t, 10) - - // Run order - - fmt.Println("1", one.GetQuota(ctx, t)) - one.GetAllowance(ctx, t) - - one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - - fmt.Println("1a", one.GetQuota(ctx, t)) - - fmt.Println("2a", two.GetQuota(ctx, t)) - two.GetAllowance(ctx, t) - - // End commit - assert.NoError(t, one.Done()) - - fmt.Println("2a", two.GetQuota(ctx, t)) - two.GetAllowance(ctx, t) - - assert.NoError(t, two.Done()) - assert.NoError(t, three.Done()) - - //allow, err = db.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ - // OwnerID: user.ID, - // OrganizationID: org.ID, - //}) - //require.NoError(t, err) - //fmt.Println(allow) - }) + // Workspace + // User: 1 + // Tpl: 1 + // Org: 1 + //myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + // OrganizationID: org.Org.ID, + // OwnerID: user.ID, + //}).Do() + // + //// Workspace + //// User: 1 + //// Tpl: 1 + //// Org: 1 + //mySecondWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + // OrganizationID: org.Org.ID, + // OwnerID: user.ID, + // TemplateID: myWorkspace.Workspace.TemplateID, + //}).Do() + // + //// Workspace + //// User: 1 + //// Tpl: 2 + //// Org: 1 + //myThirdWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + // OrganizationID: org.Org.ID, + // OwnerID: user.ID, + //}).Do() + // + //// Workspace + //// User: 1 + //// Tpl: 3 + //// Org: 2 + //myFourthWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + // OrganizationID: otherOrg.Org.ID, + // OwnerID: user.ID, + //}).Do() + // + //// Workspace + //// User: 2 + //// Tpl: 1 + //// Org: 1 + //otherWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + // OrganizationID: org.Org.ID, + // OwnerID: otherUser.ID, + //}).Do() + // + //workspaceThree := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + // OrganizationID: org.Org.ID, + // OwnerID: otherUser.ID, + // TemplateID: workspaceResp.Workspace.TemplateID, + //}).Do() + //workspaceThreeResp := dbfake.WorkspaceBuild(t, db, workspaceThree).Seed(seed).Do() + // + //ctx := testutil.Context(t, testutil.WaitLong) + //ctx = dbauthz.AsSystemRestricted(ctx) + // + //// Include all user's groups + // + //tpl := dbgen.Template(t, db, database.Template{ + // OrganizationID: org.ID, + // CreatedBy: user.ID, + //}) + //tplVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + // TemplateID: uuid.NullUUID{ + // UUID: tpl.ID, + // Valid: true, + // }, + // OrganizationID: org.ID, + // CreatedBy: user.ID, + //}) + // + //seed := database.WorkspaceBuild{ + // TemplateVersionID: tplVersion.ID, + //} + //workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + // OrganizationID: org.ID, + // OwnerID: user.ID, + // TemplateID: tpl.ID, + //}) + // + //workspaceResp := dbfake.WorkspaceBuild(t, db, workspace).Seed(seed).Do() + // + //workspaceTwo := dbgen.Workspace(t, db, database.WorkspaceTable{ + // OrganizationID: org.ID, + // OwnerID: user.ID, + // TemplateID: tpl.ID, + //}) + //workspaceTwoResp := dbfake.WorkspaceBuild(t, db, workspaceTwo).Seed(seed).Do() + // + //workspaceThree := dbgen.Workspace(t, db, database.WorkspaceTable{ + // OrganizationID: org.ID, + // OwnerID: otherUser.ID, + // TemplateID: tpl.ID, + //}) + //workspaceThreeResp := dbfake.WorkspaceBuild(t, db, workspaceThree).Seed(seed).Do() + // + //workspaceFour := dbgen.Workspace(t, db, database.WorkspaceTable{ + // OrganizationID: org.ID, + // OwnerID: otherUser.ID, + // TemplateID: tpl.ID, + //}) + //workspaceThreeResp := dbfake.WorkspaceBuild(t, db, workspaceThree).Seed(seed).Do() + // + //// TX mixing tests. **DO NOT** run these in parallel. + //// The goal here is to mess around with different ordering of + //// transactions and queries. + // + //// UpdateBuildDeadline bumps a workspace deadline while doing a quota + //// commit. + //// + //// Note: This passes if the interrupt is run before 'GetQuota()' + //// Passing orders: + //// - BeginTX -> Bump! -> GetQuota -> GetAllowance -> UpdateCost -> EndTx + //// - BeginTX -> GetQuota -> GetAllowance -> UpdateCost -> Bump! -> EndTx + //t.Run("UpdateBuildDeadline", func(t *testing.T) { + // // +------------------------------+------------------+ + // // | Begin Tx | | + // // +------------------------------+------------------+ + // // | GetQuota(user) | | + // // +------------------------------+------------------+ + // // | | BumpDeadline(w1) | + // // +------------------------------+------------------+ + // // | GetAllowance(user) | | + // // +------------------------------+------------------+ + // // | UpdateWorkspaceBuildCost(w1) | | + // // +------------------------------+------------------+ + // // | CommitTx() | | + // // +------------------------------+------------------+ + // // pq: could not serialize access due to concurrent update + // ctx := testutil.Context(t, testutil.WaitLong) + // ctx = dbauthz.AsSystemRestricted(ctx) + // + // bumpDeadline := func() { + // err = db.InTx(func(db database.Store) error { + // err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ + // Deadline: dbtime.Now(), + // MaxDeadline: dbtime.Now(), + // UpdatedAt: dbtime.Now(), + // ID: workspaceResp.Build.ID, + // }) + // return err + // }, &database.TxOptions{ + // Isolation: sql.LevelSerializable, + // }) + // assert.NoError(t, err) + // + // } + // + // // Start TX + // // Run order + // + // quota := newCommitter(t, db, workspace, workspaceResp.Build) + // quota.GetQuota(ctx, t) // Step 1 + // bumpDeadline() // Interrupt + // quota.GetAllowance(ctx, t) // Step 2 + // quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 + // // End commit + // require.NoError(t, quota.Done()) + //}) + // + //t.Run("UpdateOtherBuildDeadline", func(t *testing.T) { + // // +------------------------------+------------------+ + // // | Begin Tx | | + // // +------------------------------+------------------+ + // // | GetQuota(user) | | + // // +------------------------------+------------------+ + // // | | BumpDeadline(w2) | + // // +------------------------------+------------------+ + // // | GetAllowance(user) | | + // // +------------------------------+------------------+ + // // | UpdateWorkspaceBuildCost(w1) | | + // // +------------------------------+------------------+ + // // | CommitTx() | | + // // +------------------------------+------------------+ + // ctx := testutil.Context(t, testutil.WaitLong) + // ctx = dbauthz.AsSystemRestricted(ctx) + // + // bumpDeadline := func() { + // err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ + // Deadline: dbtime.Now(), + // MaxDeadline: dbtime.Now(), + // UpdatedAt: dbtime.Now(), + // ID: workspaceTwoResp.Build.ID, + // }) + // require.NoError(t, err) + // } + // + // // Start TX + // // Run order + // + // quota := newCommitter(t, db, workspace, workspaceResp.Build) + // quota.GetQuota(ctx, t) // Step 1 + // bumpDeadline() // Interrupt + // quota.GetAllowance(ctx, t) // Step 2 + // quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 + // // End commit + // require.NoError(t, quota.Done()) + //}) + // + //t.Run("ReadCost", func(t *testing.T) { + // ctx := testutil.Context(t, testutil.WaitLong) + // ctx = dbauthz.AsSystemRestricted(ctx) + // + // readCost := func() { + // _, err := db.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ + // OwnerID: workspace.OwnerID, + // OrganizationID: workspace.OrganizationID, + // }) + // require.NoError(t, err) + // } + // + // // Start TX + // // Run order + // + // quota := newCommitter(t, db, workspace, workspaceResp.Build) + // quota.GetQuota(ctx, t) // Step 1 + // readCost() // Interrupt + // quota.GetAllowance(ctx, t) // Step 2 + // quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 + // + // // End commit + // require.NoError(t, quota.Done()) + //}) + // + //t.Run("AutoBuild", func(t *testing.T) { + // ctx := testutil.Context(t, testutil.WaitLong) + // ctx = dbauthz.AsSystemRestricted(ctx) + // + // auto := newautobuild(t, db, workspace, workspaceResp.Build) + // quota := newCommitter(t, db, workspace, workspaceResp.Build) + // + // // Run order + // auto.DoAllReads(ctx, t) + // + // quota.GetQuota(ctx, t) // Step 1 + // auto.DoAllReads(ctx, t) + // + // quota.GetAllowance(ctx, t) // Step 2 + // auto.DoAllReads(ctx, t) + // + // quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 + // auto.DoAllReads(ctx, t) + // auto.DoAllWrites(ctx, t) + // + // // End commit + // require.NoError(t, auto.Done()) + // require.NoError(t, quota.Done()) + //}) + // + //t.Run("DoubleCommit", func(t *testing.T) { + // // +---------------------+---------------------+ + // // | W1 Quota Tx | W2 Quota Tx | + // // +---------------------+---------------------+ + // // | Begin Tx | | + // // +---------------------+---------------------+ + // // | | Begin Tx | + // // +---------------------+---------------------+ + // // | GetQuota(w1) | | + // // +---------------------+---------------------+ + // // | GetAllowance(w1) | | + // // +---------------------+---------------------+ + // // | UpdateBuildCost(w1) | | + // // +---------------------+---------------------+ + // // | | UpdateBuildCost(w2) | + // // +---------------------+---------------------+ + // // | | GetQuota(w2) | + // // +---------------------+---------------------+ + // // | | GetAllowance(w2) | + // // +---------------------+---------------------+ + // // | CommitTx() | | + // // +---------------------+---------------------+ + // // | | CommitTx() | + // // +---------------------+---------------------+ + // // pq: could not serialize access due to read/write dependencies among transactions + // ctx := testutil.Context(t, testutil.WaitLong) + // ctx = dbauthz.AsSystemRestricted(ctx) + // + // one := newCommitter(t, db, workspace, workspaceResp.Build) + // two := newCommitter(t, db, workspaceTwo, workspaceTwoResp.Build) + // + // var _, _ = one, two + // // Run order + // one.GetQuota(ctx, t) + // one.GetAllowance(ctx, t) + // + // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + // + // two.GetQuota(ctx, t) + // two.GetAllowance(ctx, t) + // two.UpdateWorkspaceBuildCostByID(ctx, t, 10) + // + // // End commit + // assert.NoError(t, one.Done()) + // assert.NoError(t, two.Done()) + //}) + // + //t.Run("DifferentUserCommit", func(t *testing.T) { + // // +---------------------+---------------------+ + // // | W1 Quota Tx | W2 Quota Tx | + // // +---------------------+---------------------+ + // // | Begin Tx | | + // // +---------------------+---------------------+ + // // | | Begin Tx | + // // +---------------------+---------------------+ + // // | GetQuota(w1) | | + // // +---------------------+---------------------+ + // // | GetAllowance(w1) | | + // // +---------------------+---------------------+ + // // | UpdateBuildCost(w1) | | + // // +---------------------+---------------------+ + // // | | UpdateBuildCost(w3) | + // // +---------------------+---------------------+ + // // | | GetQuota(w3) | + // // +---------------------+---------------------+ + // // | | GetAllowance(w3) | + // // +---------------------+---------------------+ + // // | CommitTx() | | + // // +---------------------+---------------------+ + // // | | CommitTx() | + // // +---------------------+---------------------+ + // // pq: could not serialize access due to read/write dependencies among transactions + // ctx := testutil.Context(t, testutil.WaitLong) + // ctx = dbauthz.AsSystemRestricted(ctx) + // + // one := newCommitter(t, db, workspace, workspaceResp.Build) + // three := newCommitter(t, db, workspaceTwo, workspaceThreeResp.Build) + // + // var _, _ = one, three + // // Run order + // one.GetQuota(ctx, t) + // one.GetAllowance(ctx, t) + // + // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + // + // three.GetQuota(ctx, t) + // three.GetAllowance(ctx, t) + // three.UpdateWorkspaceBuildCostByID(ctx, t, 10) + // + // // End commit + // assert.NoError(t, one.Done()) + // assert.NoError(t, three.Done()) + //}) + // + //t.Run("BumpLastUsedAt", func(t *testing.T) { + // // +---------------------+----------------------------------+ + // // | W1 Quota Tx | | + // // +---------------------+----------------------------------+ + // // | Begin Tx | | + // // +---------------------+----------------------------------+ + // // | GetQuota(w1) | | + // // +---------------------+----------------------------------+ + // // | GetAllowance(w1) | | + // // +---------------------+----------------------------------+ + // // | | UpdateWorkspaceBuildDeadline(w1) | + // // +---------------------+----------------------------------+ + // // | UpdateBuildCost(w1) | | + // // +---------------------+----------------------------------+ + // // | CommitTx() | | + // // +---------------------+----------------------------------+ + // // pq: could not serialize access due to concurrent update + // ctx := testutil.Context(t, testutil.WaitShort) + // ctx = dbauthz.AsSystemRestricted(ctx) + // + // one := newCommitter(t, db, workspace, workspaceResp.Build) + // + // // Run order + // one.GetQuota(ctx, t) + // one.GetAllowance(ctx, t) + // + // err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ + // Deadline: dbtime.Now(), + // MaxDeadline: dbtime.Now(), + // UpdatedAt: dbtime.Now(), + // ID: workspaceResp.Build.ID, + // }) + // assert.NoError(t, err) + // + // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + // + // // End commit + // assert.NoError(t, one.Done()) + //}) + // + //t.Run("ActivityBump", func(t *testing.T) { + // // +---------------------+----------------------------------+ + // // | W1 Quota Tx | | + // // +---------------------+----------------------------------+ + // // | Begin Tx | | + // // +---------------------+----------------------------------+ + // // | GetQuota(w1) | | + // // +---------------------+----------------------------------+ + // // | GetAllowance(w1) | | + // // +---------------------+----------------------------------+ + // // | | ActivityBump(w1) | + // // +---------------------+----------------------------------+ + // // | UpdateBuildCost(w1) | | + // // +---------------------+----------------------------------+ + // // | CommitTx() | | + // // +---------------------+----------------------------------+ + // // pq: could not serialize access due to concurrent update + // ctx := testutil.Context(t, testutil.WaitShort) + // ctx = dbauthz.AsSystemRestricted(ctx) + // + // one := newCommitter(t, db, workspace, workspaceResp.Build) + // + // // Run order + // one.GetQuota(ctx, t) + // one.GetAllowance(ctx, t) + // + // err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ + // NextAutostart: time.Now(), + // WorkspaceID: workspaceResp.Workspace.ID, + // }) + // + // assert.NoError(t, err) + // + // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + // + // // End commit + // assert.NoError(t, one.Done()) + //}) + // + //t.Run("UserMod", func(t *testing.T) { + // // +---------------------+----------------------------------+ + // // | W1 Quota Tx | | + // // +---------------------+----------------------------------+ + // // | Begin Tx | | + // // +---------------------+----------------------------------+ + // // | GetQuota(w1) | | + // // +---------------------+----------------------------------+ + // // | GetAllowance(w1) | | + // // +---------------------+----------------------------------+ + // // | | UpdateWorkspaceBuildDeadline(w1) | + // // +---------------------+----------------------------------+ + // // | UpdateBuildCost(w1) | | + // // +---------------------+----------------------------------+ + // // | CommitTx() | | + // // +---------------------+----------------------------------+ + // // pq: could not serialize access due to concurrent update + // ctx := testutil.Context(t, testutil.WaitShort) + // ctx = dbauthz.AsSystemRestricted(ctx) + // + // db.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{ + // OrganizationID: org.ID, + // UserID: user.ID, + // }) + // + // one := newCommitter(t, db, workspace, workspaceResp.Build) + // + // // Run order + // + // err = db.RemoveUserFromAllGroups(ctx, user.ID) + // assert.NoError(t, err) + // + // err = db.DeleteOrganizationMember(ctx, database.DeleteOrganizationMemberParams{ + // OrganizationID: org.ID, + // UserID: user.ID, + // }) + // assert.NoError(t, err) + // + // err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ + // NextAutostart: time.Now(), + // WorkspaceID: workspaceResp.Workspace.ID, + // }) + // + // one.GetQuota(ctx, t) + // one.GetAllowance(ctx, t) + // + // err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ + // NextAutostart: time.Now(), + // WorkspaceID: workspaceResp.Workspace.ID, + // }) + // + // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + // + // err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ + // NextAutostart: time.Now(), + // WorkspaceID: workspaceResp.Workspace.ID, + // }) + // + // // End commit + // assert.NoError(t, one.Done()) + //}) + // + //// TODO: Try to fail a non-repeatable read only transaction + //t.Run("ReadStale", func(t *testing.T) { + // ctx := testutil.Context(t, testutil.WaitLong) + // ctx = dbauthz.AsSystemRestricted(ctx) + // + // three := newCommitter(t, db, workspace, workspaceTwoResp.Build) + // two := newCommitter(t, db, workspaceTwo, workspaceResp.Build) + // one := newCommitter(t, db, workspace, workspaceResp.Build) + // three.UpdateWorkspaceBuildCostByID(ctx, t, 10) + // + // // Run order + // + // fmt.Println("1", one.GetQuota(ctx, t)) + // one.GetAllowance(ctx, t) + // + // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + // + // fmt.Println("1a", one.GetQuota(ctx, t)) + // + // fmt.Println("2a", two.GetQuota(ctx, t)) + // two.GetAllowance(ctx, t) + // + // // End commit + // assert.NoError(t, one.Done()) + // + // fmt.Println("2a", two.GetQuota(ctx, t)) + // two.GetAllowance(ctx, t) + // + // assert.NoError(t, two.Done()) + // assert.NoError(t, three.Done()) + // + // //allow, err = db.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ + // // OwnerID: user.ID, + // // OrganizationID: org.ID, + // //}) + // //require.NoError(t, err) + // //fmt.Println(allow) + //}) // Autobuild, then quota, then autobuild read agin in the same tx // https://www.richardstrnad.ch/posts/go-sql-how-to-get-detailed-error/ @@ -810,7 +1032,7 @@ type autobuilder struct { } func newautobuild(t *testing.T, db database.Store, workspace database.WorkspaceTable, b database.WorkspaceBuild) *autobuilder { - quotaTX := dbtestutil.StartTx(t, db, &sql.TxOptions{ + quotaTX := dbtestutil.StartTx(t, db, &database.TxOptions{ Isolation: sql.LevelRepeatableRead, }) return &autobuilder{DBTx: quotaTX, w: workspace, b: b} @@ -923,6 +1145,7 @@ func (c *autobuilder) Done() error { // committer does what the CommitQuota does, but allows // stepping through the actions in the tx and controlling the // timing. +// This is a nice wrapper to make the tests more concise. type committer struct { DBTx *dbtestutil.DBTx w database.WorkspaceTable @@ -930,7 +1153,7 @@ type committer struct { } func newCommitter(t *testing.T, db database.Store, workspace database.WorkspaceTable, build database.WorkspaceBuild) *committer { - quotaTX := dbtestutil.StartTx(t, db, &sql.TxOptions{ + quotaTX := dbtestutil.StartTx(t, db, &database.TxOptions{ Isolation: sql.LevelSerializable, ReadOnly: false, }) From 5f021ad4e3a31532512adce8228a02aebb08e536 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Mon, 28 Oct 2024 18:52:55 -0400 Subject: [PATCH 06/19] add extra test --- enterprise/coderd/workspacequota_test.go | 86 +++++++++++++----------- 1 file changed, 46 insertions(+), 40 deletions(-) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index c469cf5cb4315..48593eb8a653e 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -417,13 +417,12 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+----------------------------------+ // | GetAllowance(w1) | | // +---------------------+----------------------------------+ - // | | ActivityBump(w1) | + // | | ActivityBump(w1 | // +---------------------+----------------------------------+ // | UpdateBuildCost(w1) | | // +---------------------+----------------------------------+ // | CommitTx() | | // +---------------------+----------------------------------+ - // pq: could not serialize access due to concurrent update ctx := testutil.Context(t, testutil.WaitShort) ctx = dbauthz.AsSystemRestricted(ctx) @@ -451,6 +450,51 @@ func TestWorkspaceSerialization(t *testing.T) { assert.NoError(t, one.Done()) }) + t.Run("BumpLastUsedAt", func(t *testing.T) { + // +---------------------+----------------------------------+ + // | W1 Quota Tx | | + // +---------------------+----------------------------------+ + // | Begin Tx | | + // +---------------------+----------------------------------+ + // | GetQuota(w1) | | + // +---------------------+----------------------------------+ + // | GetAllowance(w1) | | + // +---------------------+----------------------------------+ + // | | UpdateWorkspaceBuildDeadline(w1) | + // +---------------------+----------------------------------+ + // | UpdateBuildCost(w1) | | + // +---------------------+----------------------------------+ + // | CommitTx() | | + // +---------------------+----------------------------------+ + // pq: could not serialize access due to concurrent update + ctx := testutil.Context(t, testutil.WaitShort) + ctx = dbauthz.AsSystemRestricted(ctx) + + myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: org.Org.ID, + OwnerID: user.ID, + }).Do() + + one := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) + + // Run order + one.GetQuota(ctx, t) + one.GetAllowance(ctx, t) + + err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ + Deadline: dbtime.Now(), + MaxDeadline: dbtime.Now(), + UpdatedAt: dbtime.Now(), + ID: myWorkspace.Build.ID, + }) + assert.NoError(t, err) + + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + // End commit + assert.NoError(t, one.Done()) + }) + // Workspace // User: 1 // Tpl: 1 @@ -935,44 +979,6 @@ func TestWorkspaceSerialization(t *testing.T) { // assert.NoError(t, one.Done()) //}) // - //// TODO: Try to fail a non-repeatable read only transaction - //t.Run("ReadStale", func(t *testing.T) { - // ctx := testutil.Context(t, testutil.WaitLong) - // ctx = dbauthz.AsSystemRestricted(ctx) - // - // three := newCommitter(t, db, workspace, workspaceTwoResp.Build) - // two := newCommitter(t, db, workspaceTwo, workspaceResp.Build) - // one := newCommitter(t, db, workspace, workspaceResp.Build) - // three.UpdateWorkspaceBuildCostByID(ctx, t, 10) - // - // // Run order - // - // fmt.Println("1", one.GetQuota(ctx, t)) - // one.GetAllowance(ctx, t) - // - // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - // - // fmt.Println("1a", one.GetQuota(ctx, t)) - // - // fmt.Println("2a", two.GetQuota(ctx, t)) - // two.GetAllowance(ctx, t) - // - // // End commit - // assert.NoError(t, one.Done()) - // - // fmt.Println("2a", two.GetQuota(ctx, t)) - // two.GetAllowance(ctx, t) - // - // assert.NoError(t, two.Done()) - // assert.NoError(t, three.Done()) - // - // //allow, err = db.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ - // // OwnerID: user.ID, - // // OrganizationID: org.ID, - // //}) - // //require.NoError(t, err) - // //fmt.Println(allow) - //}) // Autobuild, then quota, then autobuild read agin in the same tx // https://www.richardstrnad.ch/posts/go-sql-how-to-get-detailed-error/ From 426eeb6e24144142c3a7164815714ae58418023d Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Mon, 28 Oct 2024 19:06:31 -0400 Subject: [PATCH 07/19] delete unused code --- enterprise/coderd/workspacequota_test.go | 781 +++++------------------ 1 file changed, 173 insertions(+), 608 deletions(-) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 48593eb8a653e..98c361c1acd70 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -11,7 +11,6 @@ import ( "time" "github.com/google/uuid" - "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -407,6 +406,69 @@ func TestWorkspaceSerialization(t *testing.T) { require.NoError(t, quota.Done()) }) + // UpdateOtherBuildDeadline bumps a user's other workspace deadline + // while doing a quota commit. + t.Run("UpdateOtherBuildDeadline", func(t *testing.T) { + // +------------------------------+------------------+ + // | Begin Tx | | + // +------------------------------+------------------+ + // | GetQuota(user) | | + // +------------------------------+------------------+ + // | | BumpDeadline(w2) | + // +------------------------------+------------------+ + // | GetAllowance(user) | | + // +------------------------------+------------------+ + // | UpdateWorkspaceBuildCost(w1) | | + // +------------------------------+------------------+ + // | CommitTx() | | + // +------------------------------+------------------+ + // Works! + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: org.Org.ID, + OwnerID: user.ID, + }).Do() + + // Use the same template + otherWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: org.Org.ID, + OwnerID: user.ID, + }). + Seed(database.WorkspaceBuild{ + TemplateVersionID: myWorkspace.TemplateVersion.ID, + }). + Do() + + bumpDeadline := func() { + err := db.InTx(func(db database.Store) error { + err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ + Deadline: dbtime.Now(), + MaxDeadline: dbtime.Now(), + UpdatedAt: dbtime.Now(), + ID: otherWorkspace.Build.ID, + }) + return err + }, &database.TxOptions{ + Isolation: sql.LevelSerializable, + }) + assert.NoError(t, err) + + } + + // Start TX + // Run order + + quota := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) + quota.GetQuota(ctx, t) // Step 1 + bumpDeadline() // Interrupt + quota.GetAllowance(ctx, t) // Step 2 + quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 + // End commit + require.NoError(t, quota.Done()) + }) + t.Run("ActivityBump", func(t *testing.T) { // +---------------------+----------------------------------+ // | W1 Quota Tx | | @@ -423,13 +485,19 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+----------------------------------+ // | CommitTx() | | // +---------------------+----------------------------------+ + // pq: could not serialize access due to concurrent update ctx := testutil.Context(t, testutil.WaitShort) ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OrganizationID: org.Org.ID, OwnerID: user.ID, - }).Do() + }). + Seed(database.WorkspaceBuild{ + // Make sure the bump does something + Deadline: dbtime.Now().Add(time.Hour * -20), + }). + Do() one := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) @@ -495,490 +563,109 @@ func TestWorkspaceSerialization(t *testing.T) { assert.NoError(t, one.Done()) }) - // Workspace - // User: 1 - // Tpl: 1 - // Org: 1 - //myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ - // OrganizationID: org.Org.ID, - // OwnerID: user.ID, - //}).Do() - // - //// Workspace - //// User: 1 - //// Tpl: 1 - //// Org: 1 - //mySecondWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ - // OrganizationID: org.Org.ID, - // OwnerID: user.ID, - // TemplateID: myWorkspace.Workspace.TemplateID, - //}).Do() - // - //// Workspace - //// User: 1 - //// Tpl: 2 - //// Org: 1 - //myThirdWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ - // OrganizationID: org.Org.ID, - // OwnerID: user.ID, - //}).Do() - // - //// Workspace - //// User: 1 - //// Tpl: 3 - //// Org: 2 - //myFourthWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ - // OrganizationID: otherOrg.Org.ID, - // OwnerID: user.ID, - //}).Do() - // - //// Workspace - //// User: 2 - //// Tpl: 1 - //// Org: 1 - //otherWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ - // OrganizationID: org.Org.ID, - // OwnerID: otherUser.ID, - //}).Do() - // - //workspaceThree := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ - // OrganizationID: org.Org.ID, - // OwnerID: otherUser.ID, - // TemplateID: workspaceResp.Workspace.TemplateID, - //}).Do() - //workspaceThreeResp := dbfake.WorkspaceBuild(t, db, workspaceThree).Seed(seed).Do() - // - //ctx := testutil.Context(t, testutil.WaitLong) - //ctx = dbauthz.AsSystemRestricted(ctx) - // - //// Include all user's groups - // - //tpl := dbgen.Template(t, db, database.Template{ - // OrganizationID: org.ID, - // CreatedBy: user.ID, - //}) - //tplVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{ - // TemplateID: uuid.NullUUID{ - // UUID: tpl.ID, - // Valid: true, - // }, - // OrganizationID: org.ID, - // CreatedBy: user.ID, - //}) - // - //seed := database.WorkspaceBuild{ - // TemplateVersionID: tplVersion.ID, - //} - //workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ - // OrganizationID: org.ID, - // OwnerID: user.ID, - // TemplateID: tpl.ID, - //}) - // - //workspaceResp := dbfake.WorkspaceBuild(t, db, workspace).Seed(seed).Do() - // - //workspaceTwo := dbgen.Workspace(t, db, database.WorkspaceTable{ - // OrganizationID: org.ID, - // OwnerID: user.ID, - // TemplateID: tpl.ID, - //}) - //workspaceTwoResp := dbfake.WorkspaceBuild(t, db, workspaceTwo).Seed(seed).Do() - // - //workspaceThree := dbgen.Workspace(t, db, database.WorkspaceTable{ - // OrganizationID: org.ID, - // OwnerID: otherUser.ID, - // TemplateID: tpl.ID, - //}) - //workspaceThreeResp := dbfake.WorkspaceBuild(t, db, workspaceThree).Seed(seed).Do() - // - //workspaceFour := dbgen.Workspace(t, db, database.WorkspaceTable{ - // OrganizationID: org.ID, - // OwnerID: otherUser.ID, - // TemplateID: tpl.ID, - //}) - //workspaceThreeResp := dbfake.WorkspaceBuild(t, db, workspaceThree).Seed(seed).Do() - // - //// TX mixing tests. **DO NOT** run these in parallel. - //// The goal here is to mess around with different ordering of - //// transactions and queries. - // - //// UpdateBuildDeadline bumps a workspace deadline while doing a quota - //// commit. - //// - //// Note: This passes if the interrupt is run before 'GetQuota()' - //// Passing orders: - //// - BeginTX -> Bump! -> GetQuota -> GetAllowance -> UpdateCost -> EndTx - //// - BeginTX -> GetQuota -> GetAllowance -> UpdateCost -> Bump! -> EndTx - //t.Run("UpdateBuildDeadline", func(t *testing.T) { - // // +------------------------------+------------------+ - // // | Begin Tx | | - // // +------------------------------+------------------+ - // // | GetQuota(user) | | - // // +------------------------------+------------------+ - // // | | BumpDeadline(w1) | - // // +------------------------------+------------------+ - // // | GetAllowance(user) | | - // // +------------------------------+------------------+ - // // | UpdateWorkspaceBuildCost(w1) | | - // // +------------------------------+------------------+ - // // | CommitTx() | | - // // +------------------------------+------------------+ - // // pq: could not serialize access due to concurrent update - // ctx := testutil.Context(t, testutil.WaitLong) - // ctx = dbauthz.AsSystemRestricted(ctx) - // - // bumpDeadline := func() { - // err = db.InTx(func(db database.Store) error { - // err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ - // Deadline: dbtime.Now(), - // MaxDeadline: dbtime.Now(), - // UpdatedAt: dbtime.Now(), - // ID: workspaceResp.Build.ID, - // }) - // return err - // }, &database.TxOptions{ - // Isolation: sql.LevelSerializable, - // }) - // assert.NoError(t, err) - // - // } - // - // // Start TX - // // Run order - // - // quota := newCommitter(t, db, workspace, workspaceResp.Build) - // quota.GetQuota(ctx, t) // Step 1 - // bumpDeadline() // Interrupt - // quota.GetAllowance(ctx, t) // Step 2 - // quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 - // // End commit - // require.NoError(t, quota.Done()) - //}) - // - //t.Run("UpdateOtherBuildDeadline", func(t *testing.T) { - // // +------------------------------+------------------+ - // // | Begin Tx | | - // // +------------------------------+------------------+ - // // | GetQuota(user) | | - // // +------------------------------+------------------+ - // // | | BumpDeadline(w2) | - // // +------------------------------+------------------+ - // // | GetAllowance(user) | | - // // +------------------------------+------------------+ - // // | UpdateWorkspaceBuildCost(w1) | | - // // +------------------------------+------------------+ - // // | CommitTx() | | - // // +------------------------------+------------------+ - // ctx := testutil.Context(t, testutil.WaitLong) - // ctx = dbauthz.AsSystemRestricted(ctx) - // - // bumpDeadline := func() { - // err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ - // Deadline: dbtime.Now(), - // MaxDeadline: dbtime.Now(), - // UpdatedAt: dbtime.Now(), - // ID: workspaceTwoResp.Build.ID, - // }) - // require.NoError(t, err) - // } - // - // // Start TX - // // Run order - // - // quota := newCommitter(t, db, workspace, workspaceResp.Build) - // quota.GetQuota(ctx, t) // Step 1 - // bumpDeadline() // Interrupt - // quota.GetAllowance(ctx, t) // Step 2 - // quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 - // // End commit - // require.NoError(t, quota.Done()) - //}) - // - //t.Run("ReadCost", func(t *testing.T) { - // ctx := testutil.Context(t, testutil.WaitLong) - // ctx = dbauthz.AsSystemRestricted(ctx) - // - // readCost := func() { - // _, err := db.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{ - // OwnerID: workspace.OwnerID, - // OrganizationID: workspace.OrganizationID, - // }) - // require.NoError(t, err) - // } - // - // // Start TX - // // Run order - // - // quota := newCommitter(t, db, workspace, workspaceResp.Build) - // quota.GetQuota(ctx, t) // Step 1 - // readCost() // Interrupt - // quota.GetAllowance(ctx, t) // Step 2 - // quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 - // - // // End commit - // require.NoError(t, quota.Done()) - //}) - // - //t.Run("AutoBuild", func(t *testing.T) { - // ctx := testutil.Context(t, testutil.WaitLong) - // ctx = dbauthz.AsSystemRestricted(ctx) - // - // auto := newautobuild(t, db, workspace, workspaceResp.Build) - // quota := newCommitter(t, db, workspace, workspaceResp.Build) - // - // // Run order - // auto.DoAllReads(ctx, t) - // - // quota.GetQuota(ctx, t) // Step 1 - // auto.DoAllReads(ctx, t) - // - // quota.GetAllowance(ctx, t) // Step 2 - // auto.DoAllReads(ctx, t) - // - // quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 - // auto.DoAllReads(ctx, t) - // auto.DoAllWrites(ctx, t) - // - // // End commit - // require.NoError(t, auto.Done()) - // require.NoError(t, quota.Done()) - //}) - // - //t.Run("DoubleCommit", func(t *testing.T) { - // // +---------------------+---------------------+ - // // | W1 Quota Tx | W2 Quota Tx | - // // +---------------------+---------------------+ - // // | Begin Tx | | - // // +---------------------+---------------------+ - // // | | Begin Tx | - // // +---------------------+---------------------+ - // // | GetQuota(w1) | | - // // +---------------------+---------------------+ - // // | GetAllowance(w1) | | - // // +---------------------+---------------------+ - // // | UpdateBuildCost(w1) | | - // // +---------------------+---------------------+ - // // | | UpdateBuildCost(w2) | - // // +---------------------+---------------------+ - // // | | GetQuota(w2) | - // // +---------------------+---------------------+ - // // | | GetAllowance(w2) | - // // +---------------------+---------------------+ - // // | CommitTx() | | - // // +---------------------+---------------------+ - // // | | CommitTx() | - // // +---------------------+---------------------+ - // // pq: could not serialize access due to read/write dependencies among transactions - // ctx := testutil.Context(t, testutil.WaitLong) - // ctx = dbauthz.AsSystemRestricted(ctx) - // - // one := newCommitter(t, db, workspace, workspaceResp.Build) - // two := newCommitter(t, db, workspaceTwo, workspaceTwoResp.Build) - // - // var _, _ = one, two - // // Run order - // one.GetQuota(ctx, t) - // one.GetAllowance(ctx, t) - // - // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - // - // two.GetQuota(ctx, t) - // two.GetAllowance(ctx, t) - // two.UpdateWorkspaceBuildCostByID(ctx, t, 10) - // - // // End commit - // assert.NoError(t, one.Done()) - // assert.NoError(t, two.Done()) - //}) - // - //t.Run("DifferentUserCommit", func(t *testing.T) { - // // +---------------------+---------------------+ - // // | W1 Quota Tx | W2 Quota Tx | - // // +---------------------+---------------------+ - // // | Begin Tx | | - // // +---------------------+---------------------+ - // // | | Begin Tx | - // // +---------------------+---------------------+ - // // | GetQuota(w1) | | - // // +---------------------+---------------------+ - // // | GetAllowance(w1) | | - // // +---------------------+---------------------+ - // // | UpdateBuildCost(w1) | | - // // +---------------------+---------------------+ - // // | | UpdateBuildCost(w3) | - // // +---------------------+---------------------+ - // // | | GetQuota(w3) | - // // +---------------------+---------------------+ - // // | | GetAllowance(w3) | - // // +---------------------+---------------------+ - // // | CommitTx() | | - // // +---------------------+---------------------+ - // // | | CommitTx() | - // // +---------------------+---------------------+ - // // pq: could not serialize access due to read/write dependencies among transactions - // ctx := testutil.Context(t, testutil.WaitLong) - // ctx = dbauthz.AsSystemRestricted(ctx) - // - // one := newCommitter(t, db, workspace, workspaceResp.Build) - // three := newCommitter(t, db, workspaceTwo, workspaceThreeResp.Build) - // - // var _, _ = one, three - // // Run order - // one.GetQuota(ctx, t) - // one.GetAllowance(ctx, t) - // - // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - // - // three.GetQuota(ctx, t) - // three.GetAllowance(ctx, t) - // three.UpdateWorkspaceBuildCostByID(ctx, t, 10) - // - // // End commit - // assert.NoError(t, one.Done()) - // assert.NoError(t, three.Done()) - //}) - // - //t.Run("BumpLastUsedAt", func(t *testing.T) { - // // +---------------------+----------------------------------+ - // // | W1 Quota Tx | | - // // +---------------------+----------------------------------+ - // // | Begin Tx | | - // // +---------------------+----------------------------------+ - // // | GetQuota(w1) | | - // // +---------------------+----------------------------------+ - // // | GetAllowance(w1) | | - // // +---------------------+----------------------------------+ - // // | | UpdateWorkspaceBuildDeadline(w1) | - // // +---------------------+----------------------------------+ - // // | UpdateBuildCost(w1) | | - // // +---------------------+----------------------------------+ - // // | CommitTx() | | - // // +---------------------+----------------------------------+ - // // pq: could not serialize access due to concurrent update - // ctx := testutil.Context(t, testutil.WaitShort) - // ctx = dbauthz.AsSystemRestricted(ctx) - // - // one := newCommitter(t, db, workspace, workspaceResp.Build) - // - // // Run order - // one.GetQuota(ctx, t) - // one.GetAllowance(ctx, t) - // - // err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ - // Deadline: dbtime.Now(), - // MaxDeadline: dbtime.Now(), - // UpdatedAt: dbtime.Now(), - // ID: workspaceResp.Build.ID, - // }) - // assert.NoError(t, err) - // - // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - // - // // End commit - // assert.NoError(t, one.Done()) - //}) - // - //t.Run("ActivityBump", func(t *testing.T) { - // // +---------------------+----------------------------------+ - // // | W1 Quota Tx | | - // // +---------------------+----------------------------------+ - // // | Begin Tx | | - // // +---------------------+----------------------------------+ - // // | GetQuota(w1) | | - // // +---------------------+----------------------------------+ - // // | GetAllowance(w1) | | - // // +---------------------+----------------------------------+ - // // | | ActivityBump(w1) | - // // +---------------------+----------------------------------+ - // // | UpdateBuildCost(w1) | | - // // +---------------------+----------------------------------+ - // // | CommitTx() | | - // // +---------------------+----------------------------------+ - // // pq: could not serialize access due to concurrent update - // ctx := testutil.Context(t, testutil.WaitShort) - // ctx = dbauthz.AsSystemRestricted(ctx) - // - // one := newCommitter(t, db, workspace, workspaceResp.Build) - // - // // Run order - // one.GetQuota(ctx, t) - // one.GetAllowance(ctx, t) - // - // err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ - // NextAutostart: time.Now(), - // WorkspaceID: workspaceResp.Workspace.ID, - // }) - // - // assert.NoError(t, err) - // - // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - // - // // End commit - // assert.NoError(t, one.Done()) - //}) - // - //t.Run("UserMod", func(t *testing.T) { - // // +---------------------+----------------------------------+ - // // | W1 Quota Tx | | - // // +---------------------+----------------------------------+ - // // | Begin Tx | | - // // +---------------------+----------------------------------+ - // // | GetQuota(w1) | | - // // +---------------------+----------------------------------+ - // // | GetAllowance(w1) | | - // // +---------------------+----------------------------------+ - // // | | UpdateWorkspaceBuildDeadline(w1) | - // // +---------------------+----------------------------------+ - // // | UpdateBuildCost(w1) | | - // // +---------------------+----------------------------------+ - // // | CommitTx() | | - // // +---------------------+----------------------------------+ - // // pq: could not serialize access due to concurrent update - // ctx := testutil.Context(t, testutil.WaitShort) - // ctx = dbauthz.AsSystemRestricted(ctx) - // - // db.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{ - // OrganizationID: org.ID, - // UserID: user.ID, - // }) - // - // one := newCommitter(t, db, workspace, workspaceResp.Build) - // - // // Run order - // - // err = db.RemoveUserFromAllGroups(ctx, user.ID) - // assert.NoError(t, err) - // - // err = db.DeleteOrganizationMember(ctx, database.DeleteOrganizationMemberParams{ - // OrganizationID: org.ID, - // UserID: user.ID, - // }) - // assert.NoError(t, err) - // - // err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ - // NextAutostart: time.Now(), - // WorkspaceID: workspaceResp.Workspace.ID, - // }) - // - // one.GetQuota(ctx, t) - // one.GetAllowance(ctx, t) - // - // err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ - // NextAutostart: time.Now(), - // WorkspaceID: workspaceResp.Workspace.ID, - // }) - // - // one.UpdateWorkspaceBuildCostByID(ctx, t, 10) - // - // err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{ - // NextAutostart: time.Now(), - // WorkspaceID: workspaceResp.Workspace.ID, - // }) - // - // // End commit - // assert.NoError(t, one.Done()) - //}) - // + t.Run("UserMod", func(t *testing.T) { + // +---------------------+----------------------------------+ + // | W1 Quota Tx | | + // +---------------------+----------------------------------+ + // | Begin Tx | | + // +---------------------+----------------------------------+ + // | GetQuota(w1) | | + // +---------------------+----------------------------------+ + // | GetAllowance(w1) | | + // +---------------------+----------------------------------+ + // | | RemoveUserFromOrg | + // +---------------------+----------------------------------+ + // | UpdateBuildCost(w1) | | + // +---------------------+----------------------------------+ + // | CommitTx() | | + // +---------------------+----------------------------------+ + // Works! + ctx := testutil.Context(t, testutil.WaitShort) + ctx = dbauthz.AsSystemRestricted(ctx) + var err error + + myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: org.Org.ID, + OwnerID: user.ID, + }).Do() + + one := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) + + // Run order + + one.GetQuota(ctx, t) + one.GetAllowance(ctx, t) + + err = db.DeleteOrganizationMember(ctx, database.DeleteOrganizationMemberParams{ + OrganizationID: myWorkspace.Workspace.OrganizationID, + UserID: user.ID, + }) + assert.NoError(t, err) + + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + // End commit + assert.NoError(t, one.Done()) + }) + + // QuotaCommit 2 workspaces in different orgs. + // Workspaces do not share templates, owners, or orgs + t.Run("DoubleQuotaWorkspaces", func(t *testing.T) { + // +---------------------+---------------------+ + // | W1 Quota Tx | W2 Quota Tx | + // +---------------------+---------------------+ + // | Begin Tx | | + // +---------------------+---------------------+ + // | | Begin Tx | + // +---------------------+---------------------+ + // | GetQuota(w1) | | + // +---------------------+---------------------+ + // | GetAllowance(w1) | | + // +---------------------+---------------------+ + // | UpdateBuildCost(w1) | | + // +---------------------+---------------------+ + // | | UpdateBuildCost(w2) | + // +---------------------+---------------------+ + // | | GetQuota(w2) | + // +---------------------+---------------------+ + // | | GetAllowance(w2) | + // +---------------------+---------------------+ + // | CommitTx() | | + // +---------------------+---------------------+ + // | | CommitTx() | + // +---------------------+---------------------+ + // pq: could not serialize access due to read/write dependencies among tra + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: org.Org.ID, + OwnerID: user.ID, + }).Do() + + myOtherWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: otherOrg.Org.ID, // Different org! + OwnerID: otherUser.ID, + }).Do() + + one := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) + two := newCommitter(t, db, myOtherWorkspace.Workspace, myOtherWorkspace.Build) + + var _, _ = one, two + // Run order + one.GetQuota(ctx, t) + one.GetAllowance(ctx, t) + + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + two.GetQuota(ctx, t) + two.GetAllowance(ctx, t) + two.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + // End commit + assert.NoError(t, one.Done()) + assert.NoError(t, two.Done()) + }) // Autobuild, then quota, then autobuild read agin in the same tx // https://www.richardstrnad.ch/posts/go-sql-how-to-get-detailed-error/ @@ -1026,128 +713,6 @@ func applyWithCost(cost int32) []*proto.Response { }} } -type autobuilder struct { - DBTx *dbtestutil.DBTx - w database.WorkspaceTable - b database.WorkspaceBuild - step int - - // some stuff in mem - lastestBuild database.WorkspaceBuild - newJob database.ProvisionerJob -} - -func newautobuild(t *testing.T, db database.Store, workspace database.WorkspaceTable, b database.WorkspaceBuild) *autobuilder { - quotaTX := dbtestutil.StartTx(t, db, &database.TxOptions{ - Isolation: sql.LevelRepeatableRead, - }) - return &autobuilder{DBTx: quotaTX, w: workspace, b: b} -} - -func (c *autobuilder) DoAllWrites(ctx context.Context, t *testing.T) { - c.InsertProvisionerJob(ctx, t) - c.InsertWorkspaceBuild(ctx, t) -} - -func (c *autobuilder) InsertProvisionerJob(ctx context.Context, t *testing.T) { - now := dbtime.Now() - job, err := c.DBTx.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ - ID: uuid.New(), - CreatedAt: now, - UpdatedAt: now, - OrganizationID: c.w.OrganizationID, - InitiatorID: c.w.OwnerID, - Provisioner: database.ProvisionerTypeTerraform, - StorageMethod: database.ProvisionerStorageMethodFile, - FileID: uuid.New(), - Type: database.ProvisionerJobTypeWorkspaceBuild, - Input: []byte("{}"), - Tags: nil, - TraceMetadata: pqtype.NullRawMessage{}, - }) - require.NoError(t, err) - c.newJob = job -} - -func (c *autobuilder) InsertWorkspaceBuild(ctx context.Context, t *testing.T) { - now := dbtime.Now() - err := c.DBTx.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{ - ID: uuid.New(), - CreatedAt: now, - UpdatedAt: now, - WorkspaceID: c.w.ID, - TemplateVersionID: c.b.TemplateVersionID, - BuildNumber: c.lastestBuild.BuildNumber + 1, - Transition: database.WorkspaceTransitionStart, - InitiatorID: c.w.OwnerID, - JobID: c.newJob.ID, - ProvisionerState: nil, - Deadline: dbtime.Now().Add(time.Hour), - MaxDeadline: dbtime.Now().Add(time.Hour), - Reason: database.BuildReasonAutostart, - }) - require.NoError(t, err) -} - -func (c *autobuilder) DoAllReads(ctx context.Context, t *testing.T) { - c.GetWorkspace(ctx, t) - c.GetUser(ctx, t) - c.GetLatestWorkspaceBuildByWorkspaceID(ctx, t) - c.GetProvisionerJobByID(ctx, t) - c.GetTemplateByID(ctx, t) -} - -func (c *autobuilder) Next(ctx context.Context, t *testing.T) { - list := c.steps() - list[c.step%len(list)](ctx, t) - c.step++ -} - -func (c *autobuilder) steps() []func(context.Context, *testing.T) { - return []func(context.Context, *testing.T){ - noReturn(c.GetWorkspace), - noReturn(c.GetUser), - noReturn(c.GetLatestWorkspaceBuildByWorkspaceID), - noReturn(c.GetProvisionerJobByID), - noReturn(c.GetTemplateByID), - } -} - -func (c *autobuilder) GetWorkspace(ctx context.Context, t *testing.T) database.Workspace { - workspace, err := c.DBTx.GetWorkspaceByID(ctx, c.w.ID) - require.NoError(t, err) - return workspace -} - -func (c *autobuilder) GetUser(ctx context.Context, t *testing.T) database.User { - user, err := c.DBTx.GetUserByID(ctx, c.w.OwnerID) - require.NoError(t, err) - return user -} - -func (c *autobuilder) GetLatestWorkspaceBuildByWorkspaceID(ctx context.Context, t *testing.T) database.WorkspaceBuild { - build, err := c.DBTx.GetLatestWorkspaceBuildByWorkspaceID(ctx, c.w.ID) - require.NoError(t, err) - c.lastestBuild = build - return build -} - -func (c *autobuilder) GetProvisionerJobByID(ctx context.Context, t *testing.T) database.ProvisionerJob { - job, err := c.DBTx.GetProvisionerJobByID(ctx, c.lastestBuild.JobID) - require.NoError(t, err) - return job -} - -func (c *autobuilder) GetTemplateByID(ctx context.Context, t *testing.T) database.Template { - tpl, err := c.DBTx.GetTemplateByID(ctx, c.w.TemplateID) - require.NoError(t, err) - return tpl -} - -func (c *autobuilder) Done() error { - return c.DBTx.Done() -} - // committer does what the CommitQuota does, but allows // stepping through the actions in the tx and controlling the // timing. From b69dd0f0974a96259be62ec27a547452603c6b64 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Mon, 28 Oct 2024 19:13:06 -0400 Subject: [PATCH 08/19] revert some changes --- coderd/database/dbtestutil/tx.go | 32 ++++++++++++++++++----------- enterprise/coderd/coderd.go | 2 +- enterprise/coderd/workspacequota.go | 4 ++-- 3 files changed, 23 insertions(+), 15 deletions(-) diff --git a/coderd/database/dbtestutil/tx.go b/coderd/database/dbtestutil/tx.go index 5466657d8f9a0..6d452b41f23d9 100644 --- a/coderd/database/dbtestutil/tx.go +++ b/coderd/database/dbtestutil/tx.go @@ -10,13 +10,21 @@ import ( type DBTx struct { database.Store mu sync.Mutex - err error - errC chan error + done chan error finalErr chan error } // StartTx starts a transaction and returns a DBTx object. This allows running // 2 transactions concurrently in a test more easily. +// Example: +// +// a := StartTx(t, db, opts) +// b := StartTx(t, db, opts) +// +// a.GetUsers(...) +// b.GetUsers(...) +// +// require.NoError(t, a.Done() func StartTx(t *testing.T, db database.Store, opts *database.TxOptions) *DBTx { errC := make(chan error) finalErr := make(chan error) @@ -34,9 +42,16 @@ func StartTx(t *testing.T, db database.Store, opts *database.TxOptions) *DBTx { }) count++ if count > 1 { + // If you recursively call InTx, then don't use this. t.Logf("InTx called more than once: %d", count) + t.Fatal("InTx called more than once, this is not allowed with the StartTx helper") } - return <-errC + + select { + case _, _ = <-errC: + } + // Just return nil. The caller should be checking their own errors. + return nil }, opts) finalErr <- err }() @@ -44,13 +59,7 @@ func StartTx(t *testing.T, db database.Store, opts *database.TxOptions) *DBTx { txStore := <-txC close(txC) - return &DBTx{Store: txStore, errC: errC, finalErr: finalErr} -} - -func (tx *DBTx) SetError(err error) { - tx.mu.Lock() - defer tx.mu.Unlock() - tx.err = err + return &DBTx{Store: txStore, done: errC, finalErr: finalErr} } // Done can only be called once. If you call it twice, it will panic. @@ -58,7 +67,6 @@ func (tx *DBTx) Done() error { tx.mu.Lock() defer tx.mu.Unlock() - tx.errC <- tx.err - close(tx.errC) + close(tx.done) return <-tx.finalErr } diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 24dc43f28b576..7e59eb341411f 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -701,7 +701,7 @@ func (api *API) updateEntitlements(ctx context.Context) error { if initial, changed, enabled := featureChanged(codersdk.FeatureTemplateRBAC); shouldUpdate(initial, changed, enabled) { if enabled { - committer := Committer{ + committer := committer{ Log: api.Logger.Named("quota_committer"), Database: api.Database, } diff --git a/enterprise/coderd/workspacequota.go b/enterprise/coderd/workspacequota.go index ae9ffe14b6d89..7ea42ea24f491 100644 --- a/enterprise/coderd/workspacequota.go +++ b/enterprise/coderd/workspacequota.go @@ -18,12 +18,12 @@ import ( "github.com/coder/coder/v2/provisionerd/proto" ) -type Committer struct { +type committer struct { Log slog.Logger Database database.Store } -func (c *Committer) CommitQuota( +func (c *committer) CommitQuota( ctx context.Context, request *proto.CommitQuotaRequest, ) (*proto.CommitQuotaResponse, error) { jobID, err := uuid.Parse(request.JobId) From ec8f799148fb7c0aa3c714c53e6fee87d5a79a09 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Mon, 28 Oct 2024 19:15:21 -0400 Subject: [PATCH 09/19] fix comment --- enterprise/coderd/workspacequota_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 98c361c1acd70..261bc564de030 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -634,7 +634,7 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+---------------------+ // | | CommitTx() | // +---------------------+---------------------+ - // pq: could not serialize access due to read/write dependencies among tra + // pq: could not serialize access due to read/write dependencies among transactions ctx := testutil.Context(t, testutil.WaitLong) ctx = dbauthz.AsSystemRestricted(ctx) From f613f319fd02fd4c437047abab68d36fb279b481 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Mon, 28 Oct 2024 19:34:00 -0400 Subject: [PATCH 10/19] prevent seq scan on workspace_builds for GetQuotaAllowanceForUser --- coderd/database/queries.sql.go | 14 +++++++++----- coderd/database/queries/quotas.sql | 15 ++++++++++----- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 45cbef3f5e1d8..2ca88843e0ea0 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6736,14 +6736,18 @@ const getQuotaConsumedForUser = `-- name: GetQuotaConsumedForUser :one WITH latest_builds AS ( SELECT DISTINCT ON - (workspace_id) id, - workspace_id, - daily_cost + (wb.workspace_id) wb.id, + wb.workspace_id, + wb.daily_cost FROM workspace_builds wb +INNER JOIN + workspaces on wb.workspace_id = workspaces.id +WHERE + workspaces.owner_id = $1 ORDER BY - workspace_id, - created_at DESC + wb.workspace_id, + wb.created_at DESC ) SELECT coalesce(SUM(daily_cost), 0)::BIGINT diff --git a/coderd/database/queries/quotas.sql b/coderd/database/queries/quotas.sql index 48f9209783e4e..b96bc4c59f9a2 100644 --- a/coderd/database/queries/quotas.sql +++ b/coderd/database/queries/quotas.sql @@ -18,14 +18,19 @@ INNER JOIN groups ON WITH latest_builds AS ( SELECT DISTINCT ON - (workspace_id) id, - workspace_id, - daily_cost + (wb.workspace_id) wb.id, + wb.workspace_id, + wb.daily_cost FROM workspace_builds wb +-- This INNER JOIN prevents a +INNER JOIN + workspaces on wb.workspace_id = workspaces.id +WHERE + workspaces.owner_id = @owner_id ORDER BY - workspace_id, - created_at DESC + wb.workspace_id, + wb.created_at DESC ) SELECT coalesce(SUM(daily_cost), 0)::BIGINT From 746203d89123cd8e6823f0a9372b5cf1f69e93b1 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Mon, 28 Oct 2024 19:48:48 -0400 Subject: [PATCH 11/19] update tests --- coderd/database/querier.go | 3 + coderd/database/queries.sql.go | 10 +- coderd/database/queries/quotas.sql | 11 +- enterprise/coderd/workspacequota_test.go | 138 +++++++++++++++++++++-- 4 files changed, 144 insertions(+), 18 deletions(-) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index fcb58a7d6e305..fbf31b559ef39 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -207,6 +207,9 @@ type sqlcQuerier interface { GetProvisionerKeyByName(ctx context.Context, arg GetProvisionerKeyByNameParams) (ProvisionerKey, error) GetProvisionerLogsAfterID(ctx context.Context, arg GetProvisionerLogsAfterIDParams) ([]ProvisionerJobLog, error) GetQuotaAllowanceForUser(ctx context.Context, arg GetQuotaAllowanceForUserParams) (int64, error) + // This INNER JOIN prevents a seq scan of the workspace_builds table. + // Limit the rows to the absolute minimum required, which is all workspaces + // in a given organization for a given user. GetQuotaConsumedForUser(ctx context.Context, arg GetQuotaConsumedForUserParams) (int64, error) GetReplicaByID(ctx context.Context, id uuid.UUID) (Replica, error) GetReplicasUpdatedAfter(ctx context.Context, updatedAt time.Time) ([]Replica, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 2ca88843e0ea0..34fece1776620 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6744,7 +6744,8 @@ FROM INNER JOIN workspaces on wb.workspace_id = workspaces.id WHERE - workspaces.owner_id = $1 + workspaces.owner_id = $1 AND + workspaces.organization_id = $2 ORDER BY wb.workspace_id, wb.created_at DESC @@ -6755,8 +6756,8 @@ FROM workspaces JOIN latest_builds ON latest_builds.workspace_id = workspaces.id -WHERE NOT - deleted AND +WHERE + NOT deleted AND workspaces.owner_id = $1 AND workspaces.organization_id = $2 ` @@ -6766,6 +6767,9 @@ type GetQuotaConsumedForUserParams struct { OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` } +// This INNER JOIN prevents a seq scan of the workspace_builds table. +// Limit the rows to the absolute minimum required, which is all workspaces +// in a given organization for a given user. func (q *sqlQuerier) GetQuotaConsumedForUser(ctx context.Context, arg GetQuotaConsumedForUserParams) (int64, error) { row := q.db.QueryRowContext(ctx, getQuotaConsumedForUser, arg.OwnerID, arg.OrganizationID) var column_1 int64 diff --git a/coderd/database/queries/quotas.sql b/coderd/database/queries/quotas.sql index b96bc4c59f9a2..4a34c60bba2dc 100644 --- a/coderd/database/queries/quotas.sql +++ b/coderd/database/queries/quotas.sql @@ -23,11 +23,14 @@ SELECT wb.daily_cost FROM workspace_builds wb --- This INNER JOIN prevents a +-- This INNER JOIN prevents a seq scan of the workspace_builds table. +-- Limit the rows to the absolute minimum required, which is all workspaces +-- in a given organization for a given user. INNER JOIN workspaces on wb.workspace_id = workspaces.id WHERE - workspaces.owner_id = @owner_id + workspaces.owner_id = @owner_id AND + workspaces.organization_id = @organization_id ORDER BY wb.workspace_id, wb.created_at DESC @@ -38,8 +41,8 @@ FROM workspaces JOIN latest_builds ON latest_builds.workspace_id = workspaces.id -WHERE NOT - deleted AND +WHERE + NOT deleted AND workspaces.owner_id = @owner_id AND workspaces.organization_id = @organization_id ; diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 261bc564de030..45535b454945b 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -356,6 +356,8 @@ func TestWorkspaceSerialization(t *testing.T) { // - BeginTX -> Bump! -> GetQuota -> GetAllowance -> UpdateCost -> EndTx // - BeginTX -> GetQuota -> GetAllowance -> UpdateCost -> Bump! -> EndTx t.Run("UpdateBuildDeadline", func(t *testing.T) { + t.Skip("Expected to fail. As long as quota & deadline are on the same " + + " table and affect the same row, this will likely always fail.") // +------------------------------+------------------+ // | Begin Tx | | // +------------------------------+------------------+ @@ -470,6 +472,8 @@ func TestWorkspaceSerialization(t *testing.T) { }) t.Run("ActivityBump", func(t *testing.T) { + t.Skip("Expected to fail. As long as quota & deadline are on the same " + + " table and affect the same row, this will likely always fail.") // +---------------------+----------------------------------+ // | W1 Quota Tx | | // +---------------------+----------------------------------+ @@ -528,7 +532,7 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+----------------------------------+ // | GetAllowance(w1) | | // +---------------------+----------------------------------+ - // | | UpdateWorkspaceBuildDeadline(w1) | + // | | UpdateWorkspaceLastUsedAt(w1) | // +---------------------+----------------------------------+ // | UpdateBuildCost(w1) | | // +---------------------+----------------------------------+ @@ -549,11 +553,9 @@ func TestWorkspaceSerialization(t *testing.T) { one.GetQuota(ctx, t) one.GetAllowance(ctx, t) - err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ - Deadline: dbtime.Now(), - MaxDeadline: dbtime.Now(), - UpdatedAt: dbtime.Now(), - ID: myWorkspace.Build.ID, + err := db.UpdateWorkspaceLastUsedAt(ctx, database.UpdateWorkspaceLastUsedAtParams{ + ID: myWorkspace.Workspace.ID, + LastUsedAt: dbtime.Now(), }) assert.NoError(t, err) @@ -610,7 +612,7 @@ func TestWorkspaceSerialization(t *testing.T) { // QuotaCommit 2 workspaces in different orgs. // Workspaces do not share templates, owners, or orgs - t.Run("DoubleQuotaWorkspaces", func(t *testing.T) { + t.Run("DoubleQuotaUnrelatedWorkspaces", func(t *testing.T) { // +---------------------+---------------------+ // | W1 Quota Tx | W2 Quota Tx | // +---------------------+---------------------+ @@ -634,7 +636,6 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+---------------------+ // | | CommitTx() | // +---------------------+---------------------+ - // pq: could not serialize access due to read/write dependencies among transactions ctx := testutil.Context(t, testutil.WaitLong) ctx = dbauthz.AsSystemRestricted(ctx) @@ -667,9 +668,124 @@ func TestWorkspaceSerialization(t *testing.T) { assert.NoError(t, two.Done()) }) - // Autobuild, then quota, then autobuild read agin in the same tx - // https://www.richardstrnad.ch/posts/go-sql-how-to-get-detailed-error/ - // https://blog.danslimmon.com/2024/01/10/why-transaction-order-matters-even-if-youre-only-reading/ + // QuotaCommit 2 workspaces in different orgs. + // Workspaces do not share templates or orgs + t.Run("DoubleQuotaUserWorkspacesDiffOrgs", func(t *testing.T) { + // +---------------------+---------------------+ + // | W1 Quota Tx | W2 Quota Tx | + // +---------------------+---------------------+ + // | Begin Tx | | + // +---------------------+---------------------+ + // | | Begin Tx | + // +---------------------+---------------------+ + // | GetQuota(w1) | | + // +---------------------+---------------------+ + // | GetAllowance(w1) | | + // +---------------------+---------------------+ + // | UpdateBuildCost(w1) | | + // +---------------------+---------------------+ + // | | UpdateBuildCost(w2) | + // +---------------------+---------------------+ + // | | GetQuota(w2) | + // +---------------------+---------------------+ + // | | GetAllowance(w2) | + // +---------------------+---------------------+ + // | CommitTx() | | + // +---------------------+---------------------+ + // | | CommitTx() | + // +---------------------+---------------------+ + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: org.Org.ID, + OwnerID: user.ID, + }).Do() + + myOtherWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: otherOrg.Org.ID, // Different org! + OwnerID: user.ID, + }).Do() + + one := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) + two := newCommitter(t, db, myOtherWorkspace.Workspace, myOtherWorkspace.Build) + + var _, _ = one, two + // Run order + one.GetQuota(ctx, t) + one.GetAllowance(ctx, t) + + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + two.GetQuota(ctx, t) + two.GetAllowance(ctx, t) + two.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + // End commit + assert.NoError(t, one.Done()) + assert.NoError(t, two.Done()) + }) + + // QuotaCommit 2 workspaces in the same org. + // Workspaces do not share templates + t.Run("DoubleQuotaUserWorkspaces", func(t *testing.T) { + t.Skip("Setting a new build cost to a workspace in a org affects other " + + "workspaces in the same org. This is expected to fail.") + // +---------------------+---------------------+ + // | W1 Quota Tx | W2 Quota Tx | + // +---------------------+---------------------+ + // | Begin Tx | | + // +---------------------+---------------------+ + // | | Begin Tx | + // +---------------------+---------------------+ + // | GetQuota(w1) | | + // +---------------------+---------------------+ + // | GetAllowance(w1) | | + // +---------------------+---------------------+ + // | UpdateBuildCost(w1) | | + // +---------------------+---------------------+ + // | | UpdateBuildCost(w2) | + // +---------------------+---------------------+ + // | | GetQuota(w2) | + // +---------------------+---------------------+ + // | | GetAllowance(w2) | + // +---------------------+---------------------+ + // | CommitTx() | | + // +---------------------+---------------------+ + // | | CommitTx() | + // +---------------------+---------------------+ + // pq: could not serialize access due to read/write dependencies among transactions + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: org.Org.ID, + OwnerID: user.ID, + }).Do() + + myOtherWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: org.Org.ID, + OwnerID: user.ID, + }).Do() + + one := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) + two := newCommitter(t, db, myOtherWorkspace.Workspace, myOtherWorkspace.Build) + + var _, _ = one, two + // Run order + one.GetQuota(ctx, t) + one.GetAllowance(ctx, t) + + one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + two.GetQuota(ctx, t) + two.GetAllowance(ctx, t) + two.UpdateWorkspaceBuildCostByID(ctx, t, 10) + + // End commit + assert.NoError(t, one.Done()) + assert.NoError(t, two.Done()) + }) } func deprecatedQuotaEndpoint(ctx context.Context, client *codersdk.Client, userID string) (codersdk.WorkspaceQuota, error) { From c0d05db8a170a77632419a541f206492076ddcfe Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Mon, 28 Oct 2024 20:22:50 -0400 Subject: [PATCH 12/19] fixups --- coderd/database/db.go | 18 ++++++++++++++---- coderd/database/dbfake/builder.go | 1 + coderd/database/dbtestutil/db.go | 3 ++- coderd/database/dbtestutil/tx.go | 8 +++----- coderd/database/queries.sql.go | 3 +++ coderd/database/queries/quotas.sql | 3 +++ enterprise/coderd/workspacequota_test.go | 15 ++++++++++----- 7 files changed, 36 insertions(+), 15 deletions(-) diff --git a/coderd/database/db.go b/coderd/database/db.go index e2eeb10d6fb21..f08215dc725a1 100644 --- a/coderd/database/db.go +++ b/coderd/database/db.go @@ -48,12 +48,20 @@ type DBTX interface { GetContext(ctx context.Context, dest interface{}, query string, args ...interface{}) error } +func WithSerialRetryCount(count int) func(*sqlQuerier) { + return func(q *sqlQuerier) { + q.serialRetryCount = count + } +} + // New creates a new database store using a SQL database connection. -func New(sdb *sql.DB) Store { +func New(sdb *sql.DB, opts ...func(*sqlQuerier)) Store { dbx := sqlx.NewDb(sdb, "postgres") return &sqlQuerier{ db: dbx, sdb: dbx, + // This is an arbitrary number. + serialRetryCount: 3, } } @@ -104,6 +112,10 @@ type querier interface { type sqlQuerier struct { sdb *sqlx.DB db DBTX + + // serialRetryCount is the number of times to retry a transaction + // if it fails with a serialization error. + serialRetryCount int } func (*sqlQuerier) Wrappers() []string { @@ -143,11 +155,9 @@ func (q *sqlQuerier) InTx(function func(Store) error, txOpts *TxOptions) error { // If we are in a transaction already, the parent InTx call will handle the retry. // We do not want to duplicate those retries. if !inTx && sqlOpts.Isolation == sql.LevelSerializable { - // This is an arbitrarily chosen number. - const retryAmount = 1 var err error attempts := 0 - for attempts = 0; attempts < retryAmount; attempts++ { + for attempts = 0; attempts < q.serialRetryCount; attempts++ { txOpts.executionCount++ err = q.runTx(function, sqlOpts) if err == nil { diff --git a/coderd/database/dbfake/builder.go b/coderd/database/dbfake/builder.go index c4731c6a4693b..2731425e9c14b 100644 --- a/coderd/database/dbfake/builder.go +++ b/coderd/database/dbfake/builder.go @@ -67,6 +67,7 @@ func (b OrganizationBuilder) Do() OrganizationResponse { org := dbgen.Organization(b.t, b.db, b.seed) ctx := testutil.Context(b.t, testutil.WaitShort) + //nolint:gocritic // builder code needs perms ctx = dbauthz.AsSystemRestricted(ctx) everyone, err := b.db.InsertAllUsersGroup(ctx, org.ID) require.NoError(b.t, err) diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go index 327d880f69648..bc8c571795629 100644 --- a/coderd/database/dbtestutil/db.go +++ b/coderd/database/dbtestutil/db.go @@ -135,7 +135,8 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) { if o.dumpOnFailure { t.Cleanup(func() { DumpOnFailure(t, connectionURL) }) } - db = database.New(sqlDB) + // Unit tests should not retry serial transaction failures. + db = database.New(sqlDB, database.WithSerialRetryCount(1)) ps, err = pubsub.New(context.Background(), o.logger, sqlDB, connectionURL) require.NoError(t, err) diff --git a/coderd/database/dbtestutil/tx.go b/coderd/database/dbtestutil/tx.go index 6d452b41f23d9..a7836259c3f88 100644 --- a/coderd/database/dbtestutil/tx.go +++ b/coderd/database/dbtestutil/tx.go @@ -26,7 +26,7 @@ type DBTx struct { // // require.NoError(t, a.Done() func StartTx(t *testing.T, db database.Store, opts *database.TxOptions) *DBTx { - errC := make(chan error) + done := make(chan error) finalErr := make(chan error) txC := make(chan database.Store) @@ -47,9 +47,7 @@ func StartTx(t *testing.T, db database.Store, opts *database.TxOptions) *DBTx { t.Fatal("InTx called more than once, this is not allowed with the StartTx helper") } - select { - case _, _ = <-errC: - } + <-done // Just return nil. The caller should be checking their own errors. return nil }, opts) @@ -59,7 +57,7 @@ func StartTx(t *testing.T, db database.Store, opts *database.TxOptions) *DBTx { txStore := <-txC close(txC) - return &DBTx{Store: txStore, done: errC, finalErr: finalErr} + return &DBTx{Store: txStore, done: done, finalErr: finalErr} } // Done can only be called once. If you call it twice, it will panic. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 34fece1776620..8fd71eb3a6ca1 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6758,6 +6758,9 @@ JOIN latest_builds ON latest_builds.workspace_id = workspaces.id WHERE NOT deleted AND + -- We can likely remove these conditions since we check above. + -- But it does not hurt to be defensive and make sure future query changes + -- do not break anything. workspaces.owner_id = $1 AND workspaces.organization_id = $2 ` diff --git a/coderd/database/queries/quotas.sql b/coderd/database/queries/quotas.sql index 4a34c60bba2dc..0bc4d9a5e52d2 100644 --- a/coderd/database/queries/quotas.sql +++ b/coderd/database/queries/quotas.sql @@ -43,6 +43,9 @@ JOIN latest_builds ON latest_builds.workspace_id = workspaces.id WHERE NOT deleted AND + -- We can likely remove these conditions since we check above. + -- But it does not hurt to be defensive and make sure future query changes + -- do not break anything. workspaces.owner_id = @owner_id AND workspaces.organization_id = @organization_id ; diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 45535b454945b..a12349a0d92be 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -306,12 +306,12 @@ func TestWorkspaceQuota(t *testing.T) { }) } -// DB=ci DB_FROM=cikggwjxbths +// nolint:paralleltest // Tests must run serially func TestWorkspaceSerialization(t *testing.T) { t.Parallel() if !dbtestutil.WillUsePostgres() { - panic("We should only run this test with postgres") + t.Skip("Serialization errors only occur in postgres") } db, ps := dbtestutil.NewDB(t) @@ -342,8 +342,6 @@ func TestWorkspaceSerialization(t *testing.T) { }, user). Do() - var _ = otherOrg - // TX mixing tests. **DO NOT** run these in parallel. // The goal here is to mess around with different ordering of // transactions and queries. @@ -373,6 +371,7 @@ func TestWorkspaceSerialization(t *testing.T) { // +------------------------------+------------------+ // pq: could not serialize access due to concurrent update ctx := testutil.Context(t, testutil.WaitLong) + //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -426,6 +425,7 @@ func TestWorkspaceSerialization(t *testing.T) { // +------------------------------+------------------+ // Works! ctx := testutil.Context(t, testutil.WaitLong) + //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -456,7 +456,6 @@ func TestWorkspaceSerialization(t *testing.T) { Isolation: sql.LevelSerializable, }) assert.NoError(t, err) - } // Start TX @@ -491,6 +490,7 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+----------------------------------+ // pq: could not serialize access due to concurrent update ctx := testutil.Context(t, testutil.WaitShort) + //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -540,6 +540,7 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+----------------------------------+ // pq: could not serialize access due to concurrent update ctx := testutil.Context(t, testutil.WaitShort) + //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -583,6 +584,7 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+----------------------------------+ // Works! ctx := testutil.Context(t, testutil.WaitShort) + //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) var err error @@ -637,6 +639,7 @@ func TestWorkspaceSerialization(t *testing.T) { // | | CommitTx() | // +---------------------+---------------------+ ctx := testutil.Context(t, testutil.WaitLong) + //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -695,6 +698,7 @@ func TestWorkspaceSerialization(t *testing.T) { // | | CommitTx() | // +---------------------+---------------------+ ctx := testutil.Context(t, testutil.WaitLong) + //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -756,6 +760,7 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+---------------------+ // pq: could not serialize access due to read/write dependencies among transactions ctx := testutil.Context(t, testutil.WaitLong) + //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ From 245f6dfba3c24f14a11c4aa114c03589dd17fcd2 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Tue, 29 Oct 2024 12:34:39 -0400 Subject: [PATCH 13/19] linting --- enterprise/coderd/workspacequota_test.go | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index a12349a0d92be..c55aee5dec877 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -314,8 +314,7 @@ func TestWorkspaceSerialization(t *testing.T) { t.Skip("Serialization errors only occur in postgres") } - db, ps := dbtestutil.NewDB(t) - var _ = ps + db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) otherUser := dbgen.User(t, db, database.User{}) @@ -655,7 +654,6 @@ func TestWorkspaceSerialization(t *testing.T) { one := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) two := newCommitter(t, db, myOtherWorkspace.Workspace, myOtherWorkspace.Build) - var _, _ = one, two // Run order one.GetQuota(ctx, t) one.GetAllowance(ctx, t) @@ -714,7 +712,6 @@ func TestWorkspaceSerialization(t *testing.T) { one := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) two := newCommitter(t, db, myOtherWorkspace.Workspace, myOtherWorkspace.Build) - var _, _ = one, two // Run order one.GetQuota(ctx, t) one.GetAllowance(ctx, t) @@ -776,7 +773,6 @@ func TestWorkspaceSerialization(t *testing.T) { one := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) two := newCommitter(t, db, myOtherWorkspace.Workspace, myOtherWorkspace.Build) - var _, _ = one, two // Run order one.GetQuota(ctx, t) one.GetAllowance(ctx, t) From 422e694b567f2e039d0280decd4657b6fd7de23d Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Tue, 29 Oct 2024 12:37:02 -0400 Subject: [PATCH 14/19] fix opts --- coderd/database/db.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/coderd/database/db.go b/coderd/database/db.go index f08215dc725a1..75a28066ce905 100644 --- a/coderd/database/db.go +++ b/coderd/database/db.go @@ -57,12 +57,17 @@ func WithSerialRetryCount(count int) func(*sqlQuerier) { // New creates a new database store using a SQL database connection. func New(sdb *sql.DB, opts ...func(*sqlQuerier)) Store { dbx := sqlx.NewDb(sdb, "postgres") - return &sqlQuerier{ + q := &sqlQuerier{ db: dbx, sdb: dbx, // This is an arbitrary number. serialRetryCount: 3, } + + for _, opt := range opts { + opt(q) + } + return q } // TxOptions is used to pass some execution metadata to the callers. From 685b21e0a836b2c849eb3b0c5d2f44fc001d1496 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Tue, 29 Oct 2024 13:26:12 -0400 Subject: [PATCH 15/19] linting --- coderd/database/dbfake/builder.go | 5 +++++ enterprise/coderd/workspacequota_test.go | 7 +------ 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/coderd/database/dbfake/builder.go b/coderd/database/dbfake/builder.go index 2731425e9c14b..6803374e72445 100644 --- a/coderd/database/dbfake/builder.go +++ b/coderd/database/dbfake/builder.go @@ -39,25 +39,30 @@ type OrganizationResponse struct { } func (b OrganizationBuilder) EveryoneAllowance(allowance int) OrganizationBuilder { + //nolint: revive // returns modified struct b.allUsersAllowance = int32(allowance) return b } func (b OrganizationBuilder) Seed(seed database.Organization) OrganizationBuilder { + //nolint: revive // returns modified struct b.seed = seed return b } func (b OrganizationBuilder) Members(users ...database.User) OrganizationBuilder { for _, u := range users { + //nolint: revive // returns modified struct b.members = append(b.members, u.ID) } return b } func (b OrganizationBuilder) Group(seed database.Group, members ...database.User) OrganizationBuilder { + //nolint: revive // returns modified struct b.groups[seed] = []uuid.UUID{} for _, u := range members { + //nolint: revive // returns modified struct b.groups[seed] = append(b.groups[seed], u.ID) } return b diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index c55aee5dec877..8cc58bf047014 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -60,10 +60,6 @@ func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, org func TestWorkspaceQuota(t *testing.T) { t.Parallel() - if !dbtestutil.WillUsePostgres() { - t.Fatal("We should only run this test with postgres") - } - // This first test verifies the behavior of creating and deleting workspaces. // It also tests multi-group quota stacking and the everyone group. t.Run("CreateDelete", func(t *testing.T) { @@ -306,7 +302,7 @@ func TestWorkspaceQuota(t *testing.T) { }) } -// nolint:paralleltest // Tests must run serially +// nolint:paralleltest,tparallel // Tests must run serially func TestWorkspaceSerialization(t *testing.T) { t.Parallel() @@ -391,7 +387,6 @@ func TestWorkspaceSerialization(t *testing.T) { Isolation: sql.LevelSerializable, }) assert.NoError(t, err) - } // Start TX From 33dbd920c5962a7f184b706b14234b29dd0acc1d Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Tue, 29 Oct 2024 13:33:54 -0400 Subject: [PATCH 16/19] linting --- coderd/database/dbtestutil/tx.go | 5 ++++- enterprise/coderd/workspacequota_test.go | 6 ------ 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/coderd/database/dbtestutil/tx.go b/coderd/database/dbtestutil/tx.go index a7836259c3f88..15be63dc35aeb 100644 --- a/coderd/database/dbtestutil/tx.go +++ b/coderd/database/dbtestutil/tx.go @@ -4,6 +4,9 @@ import ( "sync" "testing" + "github.com/stretchr/testify/assert" + "golang.org/x/xerrors" + "github.com/coder/coder/v2/coderd/database" ) @@ -44,7 +47,7 @@ func StartTx(t *testing.T, db database.Store, opts *database.TxOptions) *DBTx { if count > 1 { // If you recursively call InTx, then don't use this. t.Logf("InTx called more than once: %d", count) - t.Fatal("InTx called more than once, this is not allowed with the StartTx helper") + assert.NoError(t, xerrors.New("InTx called more than once, this is not allowed with the StartTx helper")) } <-done diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 8cc58bf047014..bfb9a8c915d14 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -886,9 +886,3 @@ func (c *committer) UpdateWorkspaceBuildCostByID(ctx context.Context, t *testing func (c *committer) Done() error { return c.DBTx.Done() } - -func noReturn[T any](f func(context.Context, *testing.T) T) func(context.Context, *testing.T) { - return func(ctx context.Context, t *testing.T) { - f(ctx, t) - } -} From a99882a61b7ce03e6fbea4b84af3a6fb805852e8 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Thu, 31 Oct 2024 09:14:21 -0400 Subject: [PATCH 17/19] text cleanup --- coderd/database/querier.go | 3 --- coderd/database/queries.sql.go | 11 +++++------ coderd/database/queries/quotas.sql | 11 +++++------ enterprise/coderd/workspacequota_test.go | 5 ++--- 4 files changed, 12 insertions(+), 18 deletions(-) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index fbf31b559ef39..fcb58a7d6e305 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -207,9 +207,6 @@ type sqlcQuerier interface { GetProvisionerKeyByName(ctx context.Context, arg GetProvisionerKeyByNameParams) (ProvisionerKey, error) GetProvisionerLogsAfterID(ctx context.Context, arg GetProvisionerLogsAfterIDParams) ([]ProvisionerJobLog, error) GetQuotaAllowanceForUser(ctx context.Context, arg GetQuotaAllowanceForUserParams) (int64, error) - // This INNER JOIN prevents a seq scan of the workspace_builds table. - // Limit the rows to the absolute minimum required, which is all workspaces - // in a given organization for a given user. GetQuotaConsumedForUser(ctx context.Context, arg GetQuotaConsumedForUserParams) (int64, error) GetReplicaByID(ctx context.Context, id uuid.UUID) (Replica, error) GetReplicasUpdatedAfter(ctx context.Context, updatedAt time.Time) ([]Replica, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 8fd71eb3a6ca1..db0972debdb85 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6736,11 +6736,13 @@ const getQuotaConsumedForUser = `-- name: GetQuotaConsumedForUser :one WITH latest_builds AS ( SELECT DISTINCT ON - (wb.workspace_id) wb.id, - wb.workspace_id, + (wb.workspace_id) wb.workspace_id, wb.daily_cost FROM workspace_builds wb + -- This INNER JOIN prevents a seq scan of the workspace_builds table. + -- Limit the rows to the absolute minimum required, which is all workspaces + -- in a given organization for a given user. INNER JOIN workspaces on wb.workspace_id = workspaces.id WHERE @@ -6754,7 +6756,7 @@ SELECT coalesce(SUM(daily_cost), 0)::BIGINT FROM workspaces -JOIN latest_builds ON +INNER JOIN latest_builds ON latest_builds.workspace_id = workspaces.id WHERE NOT deleted AND @@ -6770,9 +6772,6 @@ type GetQuotaConsumedForUserParams struct { OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` } -// This INNER JOIN prevents a seq scan of the workspace_builds table. -// Limit the rows to the absolute minimum required, which is all workspaces -// in a given organization for a given user. func (q *sqlQuerier) GetQuotaConsumedForUser(ctx context.Context, arg GetQuotaConsumedForUserParams) (int64, error) { row := q.db.QueryRowContext(ctx, getQuotaConsumedForUser, arg.OwnerID, arg.OrganizationID) var column_1 int64 diff --git a/coderd/database/queries/quotas.sql b/coderd/database/queries/quotas.sql index 0bc4d9a5e52d2..7ab6189dfe8a1 100644 --- a/coderd/database/queries/quotas.sql +++ b/coderd/database/queries/quotas.sql @@ -18,14 +18,13 @@ INNER JOIN groups ON WITH latest_builds AS ( SELECT DISTINCT ON - (wb.workspace_id) wb.id, - wb.workspace_id, + (wb.workspace_id) wb.workspace_id, wb.daily_cost FROM workspace_builds wb --- This INNER JOIN prevents a seq scan of the workspace_builds table. --- Limit the rows to the absolute minimum required, which is all workspaces --- in a given organization for a given user. + -- This INNER JOIN prevents a seq scan of the workspace_builds table. + -- Limit the rows to the absolute minimum required, which is all workspaces + -- in a given organization for a given user. INNER JOIN workspaces on wb.workspace_id = workspaces.id WHERE @@ -39,7 +38,7 @@ SELECT coalesce(SUM(daily_cost), 0)::BIGINT FROM workspaces -JOIN latest_builds ON +INNER JOIN latest_builds ON latest_builds.workspace_id = workspaces.id WHERE NOT deleted AND diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index bfb9a8c915d14..68f90091882d8 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -476,7 +476,7 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+----------------------------------+ // | GetAllowance(w1) | | // +---------------------+----------------------------------+ - // | | ActivityBump(w1 | + // | | ActivityBump(w1) | // +---------------------+----------------------------------+ // | UpdateBuildCost(w1) | | // +---------------------+----------------------------------+ @@ -526,13 +526,12 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+----------------------------------+ // | GetAllowance(w1) | | // +---------------------+----------------------------------+ - // | | UpdateWorkspaceLastUsedAt(w1) | + // | | UpdateWorkspaceLastUsedAt(w1) | // +---------------------+----------------------------------+ // | UpdateBuildCost(w1) | | // +---------------------+----------------------------------+ // | CommitTx() | | // +---------------------+----------------------------------+ - // pq: could not serialize access due to concurrent update ctx := testutil.Context(t, testutil.WaitShort) //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) From 6ce8bfe37c26e2afb00893711e5a03381ea98844 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Thu, 31 Oct 2024 09:24:32 -0400 Subject: [PATCH 18/19] tests to assert errors, rather then skip the test --- enterprise/coderd/workspacequota_test.go | 32 ++++++++++++++++-------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 68f90091882d8..13142f11e5717 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -349,8 +349,9 @@ func TestWorkspaceSerialization(t *testing.T) { // - BeginTX -> Bump! -> GetQuota -> GetAllowance -> UpdateCost -> EndTx // - BeginTX -> GetQuota -> GetAllowance -> UpdateCost -> Bump! -> EndTx t.Run("UpdateBuildDeadline", func(t *testing.T) { - t.Skip("Expected to fail. As long as quota & deadline are on the same " + + t.Log("Expected to fail. As long as quota & deadline are on the same " + " table and affect the same row, this will likely always fail.") + // +------------------------------+------------------+ // | Begin Tx | | // +------------------------------+------------------+ @@ -393,12 +394,17 @@ func TestWorkspaceSerialization(t *testing.T) { // Run order quota := newCommitter(t, db, myWorkspace.Workspace, myWorkspace.Build) - quota.GetQuota(ctx, t) // Step 1 - bumpDeadline() // Interrupt - quota.GetAllowance(ctx, t) // Step 2 - quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3 + quota.GetQuota(ctx, t) // Step 1 + bumpDeadline() // Interrupt + quota.GetAllowance(ctx, t) // Step 2 + + err := quota.DBTx.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{ + ID: myWorkspace.Build.ID, + DailyCost: 10, + }) // Step 3 + require.ErrorContains(t, err, "could not serialize access due to concurrent update") // End commit - require.NoError(t, quota.Done()) + require.ErrorContains(t, quota.Done(), "failed transaction") }) // UpdateOtherBuildDeadline bumps a user's other workspace deadline @@ -465,7 +471,7 @@ func TestWorkspaceSerialization(t *testing.T) { }) t.Run("ActivityBump", func(t *testing.T) { - t.Skip("Expected to fail. As long as quota & deadline are on the same " + + t.Log("Expected to fail. As long as quota & deadline are on the same " + " table and affect the same row, this will likely always fail.") // +---------------------+----------------------------------+ // | W1 Quota Tx | | @@ -510,10 +516,14 @@ func TestWorkspaceSerialization(t *testing.T) { assert.NoError(t, err) - one.UpdateWorkspaceBuildCostByID(ctx, t, 10) + err = one.DBTx.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{ + ID: myWorkspace.Build.ID, + DailyCost: 10, + }) + require.ErrorContains(t, err, "could not serialize access due to concurrent update") // End commit - assert.NoError(t, one.Done()) + assert.ErrorContains(t, one.Done(), "failed transaction") }) t.Run("BumpLastUsedAt", func(t *testing.T) { @@ -724,7 +734,7 @@ func TestWorkspaceSerialization(t *testing.T) { // QuotaCommit 2 workspaces in the same org. // Workspaces do not share templates t.Run("DoubleQuotaUserWorkspaces", func(t *testing.T) { - t.Skip("Setting a new build cost to a workspace in a org affects other " + + t.Log("Setting a new build cost to a workspace in a org affects other " + "workspaces in the same org. This is expected to fail.") // +---------------------+---------------------+ // | W1 Quota Tx | W2 Quota Tx | @@ -779,7 +789,7 @@ func TestWorkspaceSerialization(t *testing.T) { // End commit assert.NoError(t, one.Done()) - assert.NoError(t, two.Done()) + assert.ErrorContains(t, two.Done(), "could not serialize access due to read/write dependencies among transactions") }) } From 3a9270d24f43890107b194823c5aaf93ca3c9049 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Fri, 1 Nov 2024 09:25:08 -0400 Subject: [PATCH 19/19] chore: add PGLocks query to analyze what locks are held in pg (#15308) adds `PGLocks` query for debugging what pg_locks are held during transactions. --- coderd/database/db.go | 8 ++ coderd/database/dbauthz/dbauthz.go | 4 + coderd/database/dbauthz/dbauthz_test.go | 5 +- coderd/database/dbauthz/setup_test.go | 1 + coderd/database/dbmem/dbmem.go | 4 + coderd/database/dbmetrics/querymetrics.go | 7 ++ coderd/database/dbmock/dbmock.go | 15 +++ coderd/database/pglocks.go | 119 ++++++++++++++++++++++ 8 files changed, 162 insertions(+), 1 deletion(-) create mode 100644 coderd/database/pglocks.go diff --git a/coderd/database/db.go b/coderd/database/db.go index 75a28066ce905..0f923a861efb4 100644 --- a/coderd/database/db.go +++ b/coderd/database/db.go @@ -28,6 +28,7 @@ type Store interface { wrapper Ping(ctx context.Context) (time.Duration, error) + PGLocks(ctx context.Context) (PGLocks, error) InTx(func(Store) error, *TxOptions) error } @@ -218,3 +219,10 @@ func (q *sqlQuerier) runTx(function func(Store) error, txOpts *sql.TxOptions) er } return nil } + +func safeString(s *string) string { + if s == nil { + return "" + } + return *s +} diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index ae6b307b3e7d3..9bf98aade03c4 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -603,6 +603,10 @@ func (q *querier) Ping(ctx context.Context) (time.Duration, error) { return q.db.Ping(ctx) } +func (q *querier) PGLocks(ctx context.Context) (database.PGLocks, error) { + return q.db.PGLocks(ctx) +} + // InTx runs the given function in a transaction. func (q *querier) InTx(function func(querier database.Store) error, txOpts *database.TxOptions) error { return q.db.InTx(func(tx database.Store) error { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 439cf1bdaec19..ae50309e96d66 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -152,7 +152,10 @@ func TestDBAuthzRecursive(t *testing.T) { for i := 2; i < method.Type.NumIn(); i++ { ins = append(ins, reflect.New(method.Type.In(i)).Elem()) } - if method.Name == "InTx" || method.Name == "Ping" || method.Name == "Wrappers" { + if method.Name == "InTx" || + method.Name == "Ping" || + method.Name == "Wrappers" || + method.Name == "PGLocks" { continue } // Log the name of the last method, so if there is a panic, it is diff --git a/coderd/database/dbauthz/setup_test.go b/coderd/database/dbauthz/setup_test.go index df9d551101a25..52e8dd42fea9c 100644 --- a/coderd/database/dbauthz/setup_test.go +++ b/coderd/database/dbauthz/setup_test.go @@ -34,6 +34,7 @@ var errMatchAny = xerrors.New("match any error") var skipMethods = map[string]string{ "InTx": "Not relevant", "Ping": "Not relevant", + "PGLocks": "Not relevant", "Wrappers": "Not relevant", "AcquireLock": "Not relevant", "TryAcquireLock": "Not relevant", diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go index 4f54598744dd0..e38c3e107013f 100644 --- a/coderd/database/dbmem/dbmem.go +++ b/coderd/database/dbmem/dbmem.go @@ -339,6 +339,10 @@ func (*FakeQuerier) Ping(_ context.Context) (time.Duration, error) { return 0, nil } +func (*FakeQuerier) PGLocks(_ context.Context) (database.PGLocks, error) { + return []database.PGLock{}, nil +} + func (tx *fakeTx) AcquireLock(_ context.Context, id int64) error { if _, ok := tx.FakeQuerier.locks[id]; ok { return xerrors.Errorf("cannot acquire lock %d: already held", id) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 7e74aab3b9de0..e1cfec5bac9ca 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -66,6 +66,13 @@ func (m queryMetricsStore) Ping(ctx context.Context) (time.Duration, error) { return duration, err } +func (m queryMetricsStore) PGLocks(ctx context.Context) (database.PGLocks, error) { + start := time.Now() + locks, err := m.s.PGLocks(ctx) + m.queryLatencies.WithLabelValues("PGLocks").Observe(time.Since(start).Seconds()) + return locks, err +} + func (m queryMetricsStore) InTx(f func(database.Store) error, options *database.TxOptions) error { return m.dbMetrics.InTx(f, options) } diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index ffc9ab79f777e..27b398a062051 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -4299,6 +4299,21 @@ func (mr *MockStoreMockRecorder) OrganizationMembers(arg0, arg1 any) *gomock.Cal return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OrganizationMembers", reflect.TypeOf((*MockStore)(nil).OrganizationMembers), arg0, arg1) } +// PGLocks mocks base method. +func (m *MockStore) PGLocks(arg0 context.Context) (database.PGLocks, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "PGLocks", arg0) + ret0, _ := ret[0].(database.PGLocks) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// PGLocks indicates an expected call of PGLocks. +func (mr *MockStoreMockRecorder) PGLocks(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PGLocks", reflect.TypeOf((*MockStore)(nil).PGLocks), arg0) +} + // Ping mocks base method. func (m *MockStore) Ping(arg0 context.Context) (time.Duration, error) { m.ctrl.T.Helper() diff --git a/coderd/database/pglocks.go b/coderd/database/pglocks.go new file mode 100644 index 0000000000000..85e1644b3825c --- /dev/null +++ b/coderd/database/pglocks.go @@ -0,0 +1,119 @@ +package database + +import ( + "context" + "fmt" + "reflect" + "sort" + "strings" + "time" + + "github.com/jmoiron/sqlx" + + "github.com/coder/coder/v2/coderd/util/slice" +) + +// PGLock docs see: https://www.postgresql.org/docs/current/view-pg-locks.html#VIEW-PG-LOCKS +type PGLock struct { + // LockType see: https://www.postgresql.org/docs/current/monitoring-stats.html#WAIT-EVENT-LOCK-TABLE + LockType *string `db:"locktype"` + Database *string `db:"database"` // oid + Relation *string `db:"relation"` // oid + RelationName *string `db:"relation_name"` + Page *int `db:"page"` + Tuple *int `db:"tuple"` + VirtualXID *string `db:"virtualxid"` + TransactionID *string `db:"transactionid"` // xid + ClassID *string `db:"classid"` // oid + ObjID *string `db:"objid"` // oid + ObjSubID *int `db:"objsubid"` + VirtualTransaction *string `db:"virtualtransaction"` + PID int `db:"pid"` + Mode *string `db:"mode"` + Granted bool `db:"granted"` + FastPath *bool `db:"fastpath"` + WaitStart *time.Time `db:"waitstart"` +} + +func (l PGLock) Equal(b PGLock) bool { + // Lazy, but hope this works + return reflect.DeepEqual(l, b) +} + +func (l PGLock) String() string { + granted := "granted" + if !l.Granted { + granted = "waiting" + } + var details string + switch safeString(l.LockType) { + case "relation": + details = "" + case "page": + details = fmt.Sprintf("page=%d", *l.Page) + case "tuple": + details = fmt.Sprintf("page=%d tuple=%d", *l.Page, *l.Tuple) + case "virtualxid": + details = "waiting to acquire virtual tx id lock" + default: + details = "???" + } + return fmt.Sprintf("%d-%5s [%s] %s/%s/%s: %s", + l.PID, + safeString(l.TransactionID), + granted, + safeString(l.RelationName), + safeString(l.LockType), + safeString(l.Mode), + details, + ) +} + +// PGLocks returns a list of all locks in the database currently in use. +func (q *sqlQuerier) PGLocks(ctx context.Context) (PGLocks, error) { + rows, err := q.sdb.QueryContext(ctx, ` + SELECT + relation::regclass AS relation_name, + * + FROM pg_locks; + `) + if err != nil { + return nil, err + } + + defer rows.Close() + + var locks []PGLock + err = sqlx.StructScan(rows, &locks) + if err != nil { + return nil, err + } + + return locks, err +} + +type PGLocks []PGLock + +func (l PGLocks) String() string { + // Try to group things together by relation name. + sort.Slice(l, func(i, j int) bool { + return safeString(l[i].RelationName) < safeString(l[j].RelationName) + }) + + var out strings.Builder + for i, lock := range l { + if i != 0 { + _, _ = out.WriteString("\n") + } + _, _ = out.WriteString(lock.String()) + } + return out.String() +} + +// Difference returns the difference between two sets of locks. +// This is helpful to determine what changed between the two sets. +func (l PGLocks) Difference(to PGLocks) (new PGLocks, removed PGLocks) { + return slice.SymmetricDifferenceFunc(l, to, func(a, b PGLock) bool { + return a.Equal(b) + }) +}