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>

agent/reaper/reaper_stub.go

@@ -14,6 +14,6 @@ func IsInitProcess() bool {
 	return false
 }
 
-func ForkReap(pids reap.PidCh) error {
+func ForkReap(_ reap.PidCh) error {
 	return nil
 }

coderd/database/databasefake/databasefake.go

@@ -18,33 +18,54 @@ import (
 // New returns an in-memory fake of the database.
 func New() database.Store {
 	return &fakeQuerier{
-		apiKeys:             make([]database.APIKey, 0),
-		organizationMembers: make([]database.OrganizationMember, 0),
-		organizations:       make([]database.Organization, 0),
-		users:               make([]database.User, 0),
+		mutex: &sync.RWMutex{},
+		data: &data{
+			apiKeys:             make([]database.APIKey, 0),
+			organizationMembers: make([]database.OrganizationMember, 0),
+			organizations:       make([]database.Organization, 0),
+			users:               make([]database.User, 0),
 
-		auditLogs:               make([]database.AuditLog, 0),
-		files:                   make([]database.File, 0),
-		gitSSHKey:               make([]database.GitSSHKey, 0),
-		parameterSchemas:        make([]database.ParameterSchema, 0),
-		parameterValues:         make([]database.ParameterValue, 0),
-		provisionerDaemons:      make([]database.ProvisionerDaemon, 0),
-		provisionerJobAgents:    make([]database.WorkspaceAgent, 0),
-		provisionerJobLogs:      make([]database.ProvisionerJobLog, 0),
-		provisionerJobResources: make([]database.WorkspaceResource, 0),
-		provisionerJobs:         make([]database.ProvisionerJob, 0),
-		templateVersions:        make([]database.TemplateVersion, 0),
-		templates:               make([]database.Template, 0),
-		workspaceBuilds:         make([]database.WorkspaceBuild, 0),
-		workspaceApps:           make([]database.WorkspaceApp, 0),
-		workspaces:              make([]database.Workspace, 0),
+			auditLogs:               make([]database.AuditLog, 0),
+			files:                   make([]database.File, 0),
+			gitSSHKey:               make([]database.GitSSHKey, 0),
+			parameterSchemas:        make([]database.ParameterSchema, 0),
+			parameterValues:         make([]database.ParameterValue, 0),
+			provisionerDaemons:      make([]database.ProvisionerDaemon, 0),
+			provisionerJobAgents:    make([]database.WorkspaceAgent, 0),
+			provisionerJobLogs:      make([]database.ProvisionerJobLog, 0),
+			provisionerJobResources: make([]database.WorkspaceResource, 0),
+			provisionerJobs:         make([]database.ProvisionerJob, 0),
+			templateVersions:        make([]database.TemplateVersion, 0),
+			templates:               make([]database.Template, 0),
+			workspaceBuilds:         make([]database.WorkspaceBuild, 0),
+			workspaceApps:           make([]database.WorkspaceApp, 0),
+			workspaces:              make([]database.Workspace, 0),
+		},
 	}
 }
 
+type rwMutex interface {
+	Lock()
+	RLock()
+	Unlock()
+	RUnlock()
+}
+
+// inTxMutex is a no op, since inside a transaction we are already locked.
+type inTxMutex struct{}
+
+func (inTxMutex) Lock()    {}
+func (inTxMutex) RLock()   {}
+func (inTxMutex) Unlock()  {}
+func (inTxMutex) RUnlock() {}
+
 // fakeQuerier replicates database functionality to enable quick testing.
 type fakeQuerier struct {
-	mutex sync.RWMutex
+	mutex rwMutex
+	*data
+}
 
+type data struct {
 	// Legacy tables
 	apiKeys             []database.APIKey
 	organizations       []database.Organization
@@ -73,7 +94,9 @@ type fakeQuerier struct {
 
 // InTx doesn't rollback data properly for in-memory yet.
 func (q *fakeQuerier) InTx(fn func(database.Store) error) error {
-	return fn(q)
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+	return fn(&fakeQuerier{mutex: inTxMutex{}, data: q.data})
 }
 
 func (q *fakeQuerier) AcquireProvisionerJob(_ context.Context, arg database.AcquireProvisionerJobParams) (database.ProvisionerJob, error) {

coderd/database/databasefake/databasefake_test.go

@@ -1,15 +1,67 @@
 package databasefake_test
 
 import (
+	"context"
+	"database/sql"
 	"fmt"
 	"reflect"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
 
 	"github.com/coder/coder/coderd/database"
 
 	"github.com/coder/coder/coderd/database/databasefake"
 )
 
+// test that transactions don't deadlock, and that we don't see intermediate state.
+func TestInTx(t *testing.T) {
+	t.Parallel()
+
+	uut := databasefake.New()
+
+	inTx := make(chan any)
+	queriesDone := make(chan any)
+	queriesStarted := make(chan any)
+	go func() {
+		err := uut.InTx(func(tx database.Store) error {
+			close(inTx)
+			_, err := tx.InsertOrganization(context.Background(), database.InsertOrganizationParams{
+				Name: "1",
+			})
+			assert.NoError(t, err)
+			<-queriesStarted
+			time.Sleep(5 * time.Millisecond)
+			_, err = tx.InsertOrganization(context.Background(), database.InsertOrganizationParams{
+				Name: "2",
+			})
+			assert.NoError(t, err)
+			return nil
+		})
+		assert.NoError(t, err)
+	}()
+	var nums []int
+	go func() {
+		<-inTx
+		for i := 0; i < 20; i++ {
+			orgs, err := uut.GetOrganizations(context.Background())
+			if err != nil {
+				assert.ErrorIs(t, err, sql.ErrNoRows)
+			}
+			nums = append(nums, len(orgs))
+			time.Sleep(time.Millisecond)
+		}
+		close(queriesDone)
+	}()
+	close(queriesStarted)
+	<-queriesDone
+	// ensure we never saw 1 org, only 0 or 2.
+	for i := 0; i < 20; i++ {
+		assert.NotEqual(t, 1, nums[i])
+	}
+}
+
 // TestExactMethods will ensure the fake database does not hold onto excessive
 // functions. The fake database is a manual implementation, so it is possible
 // we forget to delete functions that we remove. This unit test just ensures