Skip to content

Commit 0aa66b4

Browse files
spikecurtisspikecurtis
authored
Lock the fake database during transactions (#2478)
* Lock the fake database during transactions Signed-off-by: Spike Curtis <spike@coder.com> * Add ut for fake database transactions Signed-off-by: Spike Curtis <spike@coder.com> * fix lint Signed-off-by: Spike Curtis <spike@coder.com> * Fix lint macOS Signed-off-by: Spike Curtis <spike@coder.com>
1 parent 1455603 commit 0aa66b4

File tree

3 files changed

+100
-25
lines changed

3 files changed

+100
-25
lines changed

‎agent/reaper/reaper_stub.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,6 @@ func IsInitProcess() bool {
1414
return false
1515
}
1616

17-
func ForkReap(pids reap.PidCh) error {
17+
func ForkReap(_ reap.PidCh) error {
1818
return nil
1919
}

‎coderd/database/databasefake/databasefake.go

Lines changed: 47 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -18,33 +18,54 @@ import (
1818
// New returns an in-memory fake of the database.
1919
func New() database.Store {
2020
return &fakeQuerier{
21-
apiKeys: make([]database.APIKey, 0),
22-
organizationMembers: make([]database.OrganizationMember, 0),
23-
organizations: make([]database.Organization, 0),
24-
users: make([]database.User, 0),
25-
26-
auditLogs: make([]database.AuditLog, 0),
27-
files: make([]database.File, 0),
28-
gitSSHKey: make([]database.GitSSHKey, 0),
29-
parameterSchemas: make([]database.ParameterSchema, 0),
30-
parameterValues: make([]database.ParameterValue, 0),
31-
provisionerDaemons: make([]database.ProvisionerDaemon, 0),
32-
provisionerJobAgents: make([]database.WorkspaceAgent, 0),
33-
provisionerJobLogs: make([]database.ProvisionerJobLog, 0),
34-
provisionerJobResources: make([]database.WorkspaceResource, 0),
35-
provisionerJobs: make([]database.ProvisionerJob, 0),
36-
templateVersions: make([]database.TemplateVersion, 0),
37-
templates: make([]database.Template, 0),
38-
workspaceBuilds: make([]database.WorkspaceBuild, 0),
39-
workspaceApps: make([]database.WorkspaceApp, 0),
40-
workspaces: make([]database.Workspace, 0),
41-
}
42-
}
21+
mutex: &sync.RWMutex{},
22+
data: &data{
23+
apiKeys: make([]database.APIKey, 0),
24+
organizationMembers: make([]database.OrganizationMember, 0),
25+
organizations: make([]database.Organization, 0),
26+
users: make([]database.User, 0),
27+
28+
auditLogs: make([]database.AuditLog, 0),
29+
files: make([]database.File, 0),
30+
gitSSHKey: make([]database.GitSSHKey, 0),
31+
parameterSchemas: make([]database.ParameterSchema, 0),
32+
parameterValues: make([]database.ParameterValue, 0),
33+
provisionerDaemons: make([]database.ProvisionerDaemon, 0),
34+
provisionerJobAgents: make([]database.WorkspaceAgent, 0),
35+
provisionerJobLogs: make([]database.ProvisionerJobLog, 0),
36+
provisionerJobResources: make([]database.WorkspaceResource, 0),
37+
provisionerJobs: make([]database.ProvisionerJob, 0),
38+
templateVersions: make([]database.TemplateVersion, 0),
39+
templates: make([]database.Template, 0),
40+
workspaceBuilds: make([]database.WorkspaceBuild, 0),
41+
workspaceApps: make([]database.WorkspaceApp, 0),
42+
workspaces: make([]database.Workspace, 0),
43+
},
44+
}
45+
}
46+
47+
type rwMutex interface {
48+
Lock()
49+
RLock()
50+
Unlock()
51+
RUnlock()
52+
}
53+
54+
// inTxMutex is a no op, since inside a transaction we are already locked.
55+
type inTxMutex struct{}
56+
57+
func (inTxMutex) Lock() {}
58+
func (inTxMutex) RLock() {}
59+
func (inTxMutex) Unlock() {}
60+
func (inTxMutex) RUnlock() {}
4361

4462
// fakeQuerier replicates database functionality to enable quick testing.
4563
type fakeQuerier struct {
46-
mutex sync.RWMutex
64+
mutex rwMutex
65+
*data
66+
}
4767

68+
type data struct {
4869
// Legacy tables
4970
apiKeys []database.APIKey
5071
organizations []database.Organization
@@ -73,7 +94,9 @@ type fakeQuerier struct {
7394

7495
// InTx doesn't rollback data properly for in-memory yet.
7596
func (q *fakeQuerier) InTx(fn func(database.Store) error) error {
76-
return fn(q)
97+
q.mutex.Lock()
98+
defer q.mutex.Unlock()
99+
return fn(&fakeQuerier{mutex: inTxMutex{}, data: q.data})
77100
}
78101

79102
func (q *fakeQuerier) AcquireProvisionerJob(_ context.Context, arg database.AcquireProvisionerJobParams) (database.ProvisionerJob, error) {

‎coderd/database/databasefake/databasefake_test.go

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,67 @@
11
package databasefake_test
22

33
import (
4+
"context"
5+
"database/sql"
46
"fmt"
57
"reflect"
68
"testing"
9+
"time"
10+
11+
"github.com/stretchr/testify/assert"
712

813
"github.com/coder/coder/coderd/database"
914

1015
"github.com/coder/coder/coderd/database/databasefake"
1116
)
1217

18+
// test that transactions don't deadlock, and that we don't see intermediate state.
19+
func TestInTx(t *testing.T) {
20+
t.Parallel()
21+
22+
uut := databasefake.New()
23+
24+
inTx := make(chan any)
25+
queriesDone := make(chan any)
26+
queriesStarted := make(chan any)
27+
go func() {
28+
err := uut.InTx(func(tx database.Store) error {
29+
close(inTx)
30+
_, err := tx.InsertOrganization(context.Background(), database.InsertOrganizationParams{
31+
Name: "1",
32+
})
33+
assert.NoError(t, err)
34+
<-queriesStarted
35+
time.Sleep(5 * time.Millisecond)
36+
_, err = tx.InsertOrganization(context.Background(), database.InsertOrganizationParams{
37+
Name: "2",
38+
})
39+
assert.NoError(t, err)
40+
return nil
41+
})
42+
assert.NoError(t, err)
43+
}()
44+
var nums []int
45+
go func() {
46+
<-inTx
47+
for i := 0; i < 20; i++ {
48+
orgs, err := uut.GetOrganizations(context.Background())
49+
if err != nil {
50+
assert.ErrorIs(t, err, sql.ErrNoRows)
51+
}
52+
nums = append(nums, len(orgs))
53+
time.Sleep(time.Millisecond)
54+
}
55+
close(queriesDone)
56+
}()
57+
close(queriesStarted)
58+
<-queriesDone
59+
// ensure we never saw 1 org, only 0 or 2.
60+
for i := 0; i < 20; i++ {
61+
assert.NotEqual(t, 1, nums[i])
62+
}
63+
}
64+
1365
// TestExactMethods will ensure the fake database does not hold onto excessive
1466
// functions. The fake database is a manual implementation, so it is possible
1567
// we forget to delete functions that we remove. This unit test just ensures

0 commit comments

Comments
 (0)