Skip to content

Commit 9ccd0b2

Browse files
EmyrkEmyrk
committed
wip
1 parent 6215603 commit 9ccd0b2

File tree

1 file changed

+228
-55
lines changed

1 file changed

+228
-55
lines changed

‎enterprise/coderd/workspacequota_test.go

Lines changed: 228 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,10 @@ func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, org
6161
func TestWorkspaceQuota(t *testing.T) {
6262
t.Parallel()
6363

64+
if !dbtestutil.WillUsePostgres() {
65+
t.Fatal("We should only run this test with postgres")
66+
}
67+
6468
// This first test verifies the behavior of creating and deleting workspaces.
6569
// It also tests multi-group quota stacking and the everyone group.
6670
t.Run("CreateDelete", func(t *testing.T) {
@@ -381,13 +385,71 @@ func TestWorkspaceSerialization(t *testing.T) {
381385

382386
// UpdateBuildDeadline bumps a workspace deadline while doing a quota
383387
// commit.
384-
// pq: could not serialize access due to concurrent update
385388
//
386389
// Note: This passes if the interrupt is run before 'GetQuota()'
387390
// Passing orders:
388391
// - BeginTX -> Bump! -> GetQuota -> GetAllowance -> UpdateCost -> EndTx
389392
// - BeginTX -> GetQuota -> GetAllowance -> UpdateCost -> Bump! -> EndTx
390393
t.Run("UpdateBuildDeadline", func(t *testing.T) {
394+
// +------------------------------+------------------+
395+
// | Begin Tx | |
396+
// +------------------------------+------------------+
397+
// | GetQuota(user) | |
398+
// +------------------------------+------------------+
399+
// | | BumpDeadline(w1) |
400+
// +------------------------------+------------------+
401+
// | GetAllowance(user) | |
402+
// +------------------------------+------------------+
403+
// | UpdateWorkspaceBuildCost(w1) | |
404+
// +------------------------------+------------------+
405+
// | CommitTx() | |
406+
// +------------------------------+------------------+
407+
// pq: could not serialize access due to concurrent update
408+
ctx := testutil.Context(t, testutil.WaitLong)
409+
ctx = dbauthz.AsSystemRestricted(ctx)
410+
411+
bumpDeadline := func() {
412+
err = db.InTx(func(db database.Store) error {
413+
err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
414+
Deadline: dbtime.Now(),
415+
MaxDeadline: dbtime.Now(),
416+
UpdatedAt: dbtime.Now(),
417+
ID: workspaceResp.Build.ID,
418+
})
419+
return err
420+
}, &sql.TxOptions{
421+
Isolation: sql.LevelSerializable,
422+
})
423+
assert.NoError(t, err)
424+
425+
}
426+
427+
// Start TX
428+
// Run order
429+
430+
quota := newCommitter(t, db, workspace, workspaceResp.Build)
431+
quota.GetQuota(ctx, t) // Step 1
432+
bumpDeadline() // Interrupt
433+
quota.GetAllowance(ctx, t) // Step 2
434+
quota.UpdateWorkspaceBuildCostByID(ctx, t, 10) // Step 3
435+
// End commit
436+
require.NoError(t, quota.Done())
437+
})
438+
439+
t.Run("UpdateOtherBuildDeadline", func(t *testing.T) {
440+
// +------------------------------+------------------+
441+
// | Begin Tx | |
442+
// +------------------------------+------------------+
443+
// | GetQuota(user) | |
444+
// +------------------------------+------------------+
445+
// | | BumpDeadline(w2) |
446+
// +------------------------------+------------------+
447+
// | GetAllowance(user) | |
448+
// +------------------------------+------------------+
449+
// | UpdateWorkspaceBuildCost(w1) | |
450+
// +------------------------------+------------------+
451+
// | CommitTx() | |
452+
// +------------------------------+------------------+
391453
ctx := testutil.Context(t, testutil.WaitLong)
392454
ctx = dbauthz.AsSystemRestricted(ctx)
393455

@@ -396,7 +458,7 @@ func TestWorkspaceSerialization(t *testing.T) {
396458
Deadline: dbtime.Now(),
397459
MaxDeadline: dbtime.Now(),
398460
UpdatedAt: dbtime.Now(),
399-
ID: workspaceResp.Build.ID,
461+
ID: workspaceTwoResp.Build.ID,
400462
})
401463
require.NoError(t, err)
402464
}
@@ -464,6 +526,30 @@ func TestWorkspaceSerialization(t *testing.T) {
464526
})
465527

466528
t.Run("DoubleCommit", func(t *testing.T) {
529+
// +---------------------+---------------------+
530+
// | W1 Quota Tx | W2 Quota Tx |
531+
// +---------------------+---------------------+
532+
// | Begin Tx | |
533+
// +---------------------+---------------------+
534+
// | | Begin Tx |
535+
// +---------------------+---------------------+
536+
// | GetQuota(w1) | |
537+
// +---------------------+---------------------+
538+
// | GetAllowance(w1) | |
539+
// +---------------------+---------------------+
540+
// | UpdateBuildCost(w1) | |
541+
// +---------------------+---------------------+
542+
// | | UpdateBuildCost(w2) |
543+
// +---------------------+---------------------+
544+
// | | GetQuota(w2) |
545+
// +---------------------+---------------------+
546+
// | | GetAllowance(w2) |
547+
// +---------------------+---------------------+
548+
// | CommitTx() | |
549+
// +---------------------+---------------------+
550+
// | | CommitTx() |
551+
// +---------------------+---------------------+
552+
// pq: could not serialize access due to read/write dependencies among transactions
467553
ctx := testutil.Context(t, testutil.WaitLong)
468554
ctx = dbauthz.AsSystemRestricted(ctx)
469555

@@ -476,82 +562,155 @@ func TestWorkspaceSerialization(t *testing.T) {
476562
one.GetAllowance(ctx, t)
477563

478564
one.UpdateWorkspaceBuildCostByID(ctx, t, 10)
479-
two.UpdateWorkspaceBuildCostByID(ctx, t, 10)
480565

481566
two.GetQuota(ctx, t)
482567
two.GetAllowance(ctx, t)
568+
two.UpdateWorkspaceBuildCostByID(ctx, t, 10)
483569

484570
// End commit
485571
require.NoError(t, one.Done())
486572
require.NoError(t, two.Done())
487573
})
488574

489575
t.Run("BumpLastUsedAt", func(t *testing.T) {
576+
// +---------------------+----------------------------------+
577+
// | W1 Quota Tx | |
578+
// +---------------------+----------------------------------+
579+
// | Begin Tx | |
580+
// +---------------------+----------------------------------+
581+
// | GetQuota(w1) | |
582+
// +---------------------+----------------------------------+
583+
// | GetAllowance(w1) | |
584+
// +---------------------+----------------------------------+
585+
// | | UpdateWorkspaceBuildDeadline(w1) |
586+
// +---------------------+----------------------------------+
587+
// | UpdateBuildCost(w1) | |
588+
// +---------------------+----------------------------------+
589+
// | CommitTx() | |
590+
// +---------------------+----------------------------------+
591+
// pq: could not serialize access due to concurrent update
490592
ctx := testutil.Context(t, testutil.WaitShort)
491593
ctx = dbauthz.AsSystemRestricted(ctx)
492594

493-
two := dbtestutil.StartTx(t, db, &sql.TxOptions{
494-
Isolation: sql.LevelSerializable,
495-
ReadOnly: false,
496-
})
497595
one := newCommitter(t, db, workspace, workspaceResp.Build)
498596

499-
//two := newCommitter(t, db, workspaceTwo, workspaceResp.Build)
500-
501597
// Run order
502598
one.GetQuota(ctx, t)
503599
one.GetAllowance(ctx, t)
504600

505-
//two.UpdateWorkspaceBuildCostByID(ctx, t, 10)
601+
err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
602+
Deadline: dbtime.Now(),
603+
MaxDeadline: dbtime.Now(),
604+
UpdatedAt: dbtime.Now(),
605+
ID: workspaceResp.Build.ID,
606+
})
607+
assert.NoError(t, err)
506608

507-
//err := two.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{
508-
// ID: workspaceResp.Build.ID,
509-
// DailyCost: 30,
510-
//})
511-
//require.NoError(t, err)
609+
one.UpdateWorkspaceBuildCostByID(ctx, t, 10)
512610

513-
//q, err := two.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{
514-
// OwnerID: user.ID,
515-
// OrganizationID: workspace.OrganizationID,
516-
//})
517-
//require.NoError(t, err)
518-
//fmt.Println(q)
611+
// End commit
612+
assert.NoError(t, one.Done())
613+
})
614+
615+
t.Run("ActivityBump", func(t *testing.T) {
616+
// +---------------------+----------------------------------+
617+
// | W1 Quota Tx | |
618+
// +---------------------+----------------------------------+
619+
// | Begin Tx | |
620+
// +---------------------+----------------------------------+
621+
// | GetQuota(w1) | |
622+
// +---------------------+----------------------------------+
623+
// | GetAllowance(w1) | |
624+
// +---------------------+----------------------------------+
625+
// | | ActivityBump(w1) |
626+
// +---------------------+----------------------------------+
627+
// | UpdateBuildCost(w1) | |
628+
// +---------------------+----------------------------------+
629+
// | CommitTx() | |
630+
// +---------------------+----------------------------------+
631+
// pq: could not serialize access due to concurrent update
632+
ctx := testutil.Context(t, testutil.WaitShort)
633+
ctx = dbauthz.AsSystemRestricted(ctx)
634+
635+
one := newCommitter(t, db, workspace, workspaceResp.Build)
636+
637+
// Run order
638+
one.GetQuota(ctx, t)
639+
one.GetAllowance(ctx, t)
640+
641+
err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{
642+
NextAutostart: time.Now(),
643+
WorkspaceID: workspaceResp.Workspace.ID,
644+
})
645+
646+
assert.NoError(t, err)
519647

520648
one.UpdateWorkspaceBuildCostByID(ctx, t, 10)
521649

522-
wg := sync.WaitGroup{}
523-
wg.Add(1)
524-
go func() {
525-
defer wg.Done()
526-
err = two.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
527-
Deadline: dbtime.Now(),
528-
MaxDeadline: dbtime.Now(),
529-
UpdatedAt: dbtime.Now(),
530-
ID: workspaceResp.Build.ID,
531-
})
532-
require.NoError(t, err)
533-
}()
534-
time.Sleep(time.Millisecond * 800)
650+
// End commit
651+
assert.NoError(t, one.Done())
652+
})
535653

536-
//err = db.UpdateWorkspaceLastUsedAt(ctx, database.UpdateWorkspaceLastUsedAtParams{
537-
// ID: workspace.ID,
538-
// LastUsedAt: dbtime.Now(),
539-
//})
540-
//require.NoError(t, err)
541-
//
542-
//err = db.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{
543-
// ID: workspaceResp.Build.ID,
544-
// DailyCost: 20,
545-
//})
546-
//require.NoError(t, err)
654+
t.Run("UserMod", func(t *testing.T) {
655+
// +---------------------+----------------------------------+
656+
// | W1 Quota Tx | |
657+
// +---------------------+----------------------------------+
658+
// | Begin Tx | |
659+
// +---------------------+----------------------------------+
660+
// | GetQuota(w1) | |
661+
// +---------------------+----------------------------------+
662+
// | GetAllowance(w1) | |
663+
// +---------------------+----------------------------------+
664+
// | | UpdateWorkspaceBuildDeadline(w1) |
665+
// +---------------------+----------------------------------+
666+
// | UpdateBuildCost(w1) | |
667+
// +---------------------+----------------------------------+
668+
// | CommitTx() | |
669+
// +---------------------+----------------------------------+
670+
// pq: could not serialize access due to concurrent update
671+
ctx := testutil.Context(t, testutil.WaitShort)
672+
ctx = dbauthz.AsSystemRestricted(ctx)
673+
674+
db.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{
675+
OrganizationID: org.ID,
676+
UserID: user.ID,
677+
})
678+
679+
one := newCommitter(t, db, workspace, workspaceResp.Build)
680+
681+
// Run order
682+
683+
err = db.RemoveUserFromAllGroups(ctx, user.ID)
684+
assert.NoError(t, err)
685+
686+
err = db.DeleteOrganizationMember(ctx, database.DeleteOrganizationMemberParams{
687+
OrganizationID: org.ID,
688+
UserID: user.ID,
689+
})
690+
assert.NoError(t, err)
691+
692+
err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{
693+
NextAutostart: time.Now(),
694+
WorkspaceID: workspaceResp.Workspace.ID,
695+
})
696+
697+
one.GetQuota(ctx, t)
698+
one.GetAllowance(ctx, t)
547699

548-
//two.GetQuota(ctx, t)
549-
//two.GetAllowance(ctx, t)
700+
err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{
701+
NextAutostart: time.Now(),
702+
WorkspaceID: workspaceResp.Workspace.ID,
703+
})
704+
705+
one.UpdateWorkspaceBuildCostByID(ctx, t, 10)
706+
707+
err = db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{
708+
NextAutostart: time.Now(),
709+
WorkspaceID: workspaceResp.Workspace.ID,
710+
})
550711

551712
// End commit
552-
require.NoError(t, one.Done())
553-
wg.Wait()
554-
require.NoError(t, two.Done())
713+
assert.NoError(t, one.Done())
555714
})
556715

557716
// TODO: Try to fail a non-repeatable read only transaction
@@ -577,13 +736,13 @@ func TestWorkspaceSerialization(t *testing.T) {
577736
two.GetAllowance(ctx, t)
578737

579738
// End commit
580-
require.NoError(t, one.Done())
739+
assert.NoError(t, one.Done())
581740

582741
fmt.Println("2a", two.GetQuota(ctx, t))
583742
two.GetAllowance(ctx, t)
584743

585-
require.NoError(t, two.Done())
586-
require.NoError(t, three.Done())
744+
assert.NoError(t, two.Done())
745+
assert.NoError(t, three.Done())
587746

588747
//allow, err = db.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{
589748
// OwnerID: user.ID,
@@ -778,7 +937,12 @@ func newCommitter(t *testing.T, db database.Store, workspace database.WorkspaceT
778937
return &committer{DBTx: quotaTX, w: workspace, b: build}
779938
}
780939

940+
// GetQuota touches:
941+
// - workspace_builds
942+
// - workspaces
781943
func (c *committer) GetQuota(ctx context.Context, t *testing.T) int64 {
944+
t.Helper()
945+
782946
consumed, err := c.DBTx.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{
783947
OwnerID: c.w.OwnerID,
784948
OrganizationID: c.w.OrganizationID,
@@ -787,7 +951,14 @@ func (c *committer) GetQuota(ctx context.Context, t *testing.T) int64 {
787951
return consumed
788952
}
789953

954+
// GetAllowance touches:
955+
// - group_members_expanded
956+
// - users
957+
// - groups
958+
// - org_members
790959
func (c *committer) GetAllowance(ctx context.Context, t *testing.T) int64 {
960+
t.Helper()
961+
791962
allowance, err := c.DBTx.GetQuotaAllowanceForUser(ctx, database.GetQuotaAllowanceForUserParams{
792963
UserID: c.w.OwnerID,
793964
OrganizationID: c.w.OrganizationID,
@@ -796,12 +967,14 @@ func (c *committer) GetAllowance(ctx context.Context, t *testing.T) int64 {
796967
return allowance
797968
}
798969

799-
func (c *committer) UpdateWorkspaceBuildCostByID(ctx context.Context, t *testing.T, cost int32) {
970+
func (c *committer) UpdateWorkspaceBuildCostByID(ctx context.Context, t *testing.T, cost int32) bool {
971+
t.Helper()
972+
800973
err := c.DBTx.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{
801974
ID: c.b.ID,
802975
DailyCost: cost,
803976
})
804-
require.NoError(t, err)
977+
return assert.NoError(t, err)
805978
}
806979

807980
func (c *committer) Done() error {

0 commit comments

Comments
 (0)