Skip to content

Commit 37c71f1

Browse files
EmyrkEmyrk
committed
chore: use database in current context for file cache
Using the db.Store when in a TX causes a deadlock for dbmem. In production, this can cause a deadlock if at the current conn pool limit.
1 parent 659b787 commit 37c71f1

File tree

5 files changed

+62
-63
lines changed

5 files changed

+62
-63
lines changed

‎coderd/coderd.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -574,7 +574,7 @@ func New(options *Options) *API {
574574
TemplateScheduleStore: options.TemplateScheduleStore,
575575
UserQuietHoursScheduleStore: options.UserQuietHoursScheduleStore,
576576
AccessControlStore: options.AccessControlStore,
577-
FileCache: files.NewFromStore(options.Database, options.PrometheusRegistry, options.Authorizer),
577+
FileCache: files.New(options.PrometheusRegistry, options.Authorizer),
578578
Experiments: experiments,
579579
WebpushDispatcher: options.WebPushDispatcher,
580580
healthCheckGroup: &singleflight.Group[string, *healthsdk.HealthcheckReport]{},

‎coderd/dynamicparameters/render.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -169,14 +169,14 @@ func (r *loader) dynamicRenderer(ctx context.Context, db database.Store, cache *
169169
var templateFS fs.FS
170170
var err error
171171

172-
templateFS, err = cache.Acquire(fileCtx, r.job.FileID)
172+
templateFS, err = cache.Acquire(fileCtx, db, r.job.FileID)
173173
if err != nil {
174174
return nil, xerrors.Errorf("acquire template file: %w", err)
175175
}
176176

177177
var moduleFilesFS *files.CloseFS
178178
if r.terraformValues.CachedModuleFiles.Valid {
179-
moduleFilesFS, err = cache.Acquire(fileCtx, r.terraformValues.CachedModuleFiles.UUID)
179+
moduleFilesFS, err = cache.Acquire(fileCtx, db, r.terraformValues.CachedModuleFiles.UUID)
180180
if err != nil {
181181
return nil, xerrors.Errorf("acquire module files: %w", err)
182182
}

‎coderd/files/cache.go‎

Lines changed: 29 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -20,38 +20,15 @@ import (
2020
)
2121

2222
type FileAcquirer interface {
23-
Acquire(ctx context.Context, fileID uuid.UUID) (*CloseFS, error)
23+
Acquire(ctx context.Context, db database.Store, fileID uuid.UUID) (*CloseFS, error)
2424
}
2525

26-
// NewFromStore returns a file cache that will fetch files from the provided
27-
// database.
28-
func NewFromStore(store database.Store, registerer prometheus.Registerer, authz rbac.Authorizer) *Cache {
29-
fetch := func(ctx context.Context, fileID uuid.UUID) (CacheEntryValue, error) {
30-
// Make sure the read does not fail due to authorization issues.
31-
// Authz is checked on the Acquire call, so this is safe.
32-
//nolint:gocritic
33-
file, err := store.GetFileByID(dbauthz.AsFileReader(ctx), fileID)
34-
if err != nil {
35-
return CacheEntryValue{}, xerrors.Errorf("failed to read file from database: %w", err)
36-
}
37-
38-
content := bytes.NewBuffer(file.Data)
39-
return CacheEntryValue{
40-
Object: file.RBACObject(),
41-
FS: archivefs.FromTarReader(content),
42-
Size: int64(len(file.Data)),
43-
}, nil
44-
}
45-
46-
return New(fetch, registerer, authz)
47-
}
48-
49-
func New(fetch fetcher, registerer prometheus.Registerer, authz rbac.Authorizer) *Cache {
26+
// New returns a file cache that will fetch files from a database
27+
func New(registerer prometheus.Registerer, authz rbac.Authorizer) *Cache {
5028
return (&Cache{
51-
lock: sync.Mutex{},
52-
data: make(map[uuid.UUID]*cacheEntry),
53-
fetcher: fetch,
54-
authz: authz,
29+
lock: sync.Mutex{},
30+
data: make(map[uuid.UUID]*cacheEntry),
31+
authz: authz,
5532
}).registerMetrics(registerer)
5633
}
5734

@@ -110,9 +87,8 @@ func (c *Cache) registerMetrics(registerer prometheus.Registerer) *Cache {
11087
// loaded into memory exactly once. We hold those files until there are no
11188
// longer any open connections, and then we remove the value from the map.
11289
type Cache struct {
113-
lock sync.Mutex
114-
data map[uuid.UUID]*cacheEntry
115-
fetcher
90+
lock sync.Mutex
91+
data map[uuid.UUID]*cacheEntry
11692
authz rbac.Authorizer
11793

11894
// metrics
@@ -142,8 +118,6 @@ type cacheEntry struct {
142118
value *lazy.ValueWithError[CacheEntryValue]
143119
}
144120

145-
type fetcher func(context.Context, uuid.UUID) (CacheEntryValue, error)
146-
147121
var _ fs.FS = (*CloseFS)(nil)
148122

149123
// CloseFS is a wrapper around fs.FS that implements io.Closer. The Close()
@@ -163,12 +137,12 @@ func (f *CloseFS) Close() { f.close() }
163137
//
164138
// Safety: Every call to Acquire that does not return an error must have a
165139
// matching call to Release.
166-
func (c *Cache) Acquire(ctx context.Context, fileID uuid.UUID) (*CloseFS, error) {
140+
func (c *Cache) Acquire(ctx context.Context, db database.Store, fileID uuid.UUID) (*CloseFS, error) {
167141
// It's important that this `Load` call occurs outside `prepare`, after the
168142
// mutex has been released, or we would continue to hold the lock until the
169143
// entire file has been fetched, which may be slow, and would prevent other
170144
// files from being fetched in parallel.
171-
it, err := c.prepare(ctx, fileID).Load()
145+
it, err := c.prepare(ctx, db, fileID).Load()
172146
if err != nil {
173147
c.release(fileID)
174148
return nil, err
@@ -195,14 +169,14 @@ func (c *Cache) Acquire(ctx context.Context, fileID uuid.UUID) (*CloseFS, error)
195169
}, nil
196170
}
197171

198-
func (c *Cache) prepare(ctx context.Context, fileID uuid.UUID) *lazy.ValueWithError[CacheEntryValue] {
172+
func (c *Cache) prepare(ctx context.Context, db database.Store, fileID uuid.UUID) *lazy.ValueWithError[CacheEntryValue] {
199173
c.lock.Lock()
200174
defer c.lock.Unlock()
201175

202176
entry, ok := c.data[fileID]
203177
if !ok {
204178
value := lazy.NewWithError(func() (CacheEntryValue, error) {
205-
val, err := c.fetcher(ctx, fileID)
179+
val, err := fetch(ctx, db, fileID)
206180

207181
// Always add to the cache size the bytes of the file loaded.
208182
if err == nil {
@@ -269,3 +243,20 @@ func (c *Cache) Count() int {
269243

270244
return len(c.data)
271245
}
246+
247+
func fetch(ctx context.Context, store database.Store, fileID uuid.UUID) (CacheEntryValue, error) {
248+
// Make sure the read does not fail due to authorization issues.
249+
// Authz is checked on the Acquire call, so this is safe.
250+
//nolint:gocritic
251+
file, err := store.GetFileByID(dbauthz.AsFileReader(ctx), fileID)
252+
if err != nil {
253+
return CacheEntryValue{}, xerrors.Errorf("failed to read file from database: %w", err)
254+
}
255+
256+
content := bytes.NewBuffer(file.Data)
257+
return CacheEntryValue{
258+
Object: file.RBACObject(),
259+
FS: archivefs.FromTarReader(content),
260+
Size: int64(len(file.Data)),
261+
}, nil
262+
}

‎coderd/files/cache_test.go‎

Lines changed: 26 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@ import (
88

99
"github.com/google/uuid"
1010
"github.com/prometheus/client_golang/prometheus"
11-
"github.com/spf13/afero"
1211
"github.com/stretchr/testify/require"
12+
"go.uber.org/mock/gomock"
1313
"golang.org/x/sync/errgroup"
1414

1515
"cdr.dev/slog/sloggers/slogtest"
@@ -18,6 +18,7 @@ import (
1818
"github.com/coder/coder/v2/coderd/database"
1919
"github.com/coder/coder/v2/coderd/database/dbauthz"
2020
"github.com/coder/coder/v2/coderd/database/dbgen"
21+
"github.com/coder/coder/v2/coderd/database/dbmock"
2122
"github.com/coder/coder/v2/coderd/database/dbtestutil"
2223
"github.com/coder/coder/v2/coderd/files"
2324
"github.com/coder/coder/v2/coderd/rbac"
@@ -58,7 +59,7 @@ func TestCacheRBAC(t *testing.T) {
5859
require.Equal(t, 0, cache.Count())
5960
rec.Reset()
6061

61-
_, err := cache.Acquire(nobody, file.ID)
62+
_, err := cache.Acquire(nobody, db, file.ID)
6263
require.Error(t, err)
6364
require.True(t, rbac.IsUnauthorizedError(err))
6465

@@ -75,18 +76,18 @@ func TestCacheRBAC(t *testing.T) {
7576
require.Equal(t, 0, cache.Count())
7677

7778
// Read the file with a file reader to put it into the cache.
78-
a, err := cache.Acquire(cacheReader, file.ID)
79+
a, err := cache.Acquire(cacheReader, db, file.ID)
7980
require.NoError(t, err)
8081
require.Equal(t, 1, cache.Count())
8182

8283
// "nobody" should not be able to read the file.
83-
_, err = cache.Acquire(nobody, file.ID)
84+
_, err = cache.Acquire(nobody, db, file.ID)
8485
require.Error(t, err)
8586
require.True(t, rbac.IsUnauthorizedError(err))
8687
require.Equal(t, 1, cache.Count())
8788

8889
// UserReader can
89-
b, err := cache.Acquire(userReader, file.ID)
90+
b, err := cache.Acquire(userReader, db, file.ID)
9091
require.NoError(t, err)
9192
require.Equal(t, 1, cache.Count())
9293

@@ -110,16 +111,21 @@ func TestConcurrency(t *testing.T) {
110111
ctx := dbauthz.AsFileReader(t.Context())
111112

112113
const fileSize = 10
113-
emptyFS := afero.NewIOFS(afero.NewReadOnlyFs(afero.NewMemMapFs()))
114114
var fetches atomic.Int64
115115
reg := prometheus.NewRegistry()
116-
c := files.New(func(_ context.Context, _ uuid.UUID) (files.CacheEntryValue, error) {
116+
117+
dbM := dbmock.NewMockStore(gomock.NewController(t))
118+
dbM.EXPECT().GetFileByID(gomock.Any(), gomock.Any()).DoAndReturn(func(mTx context.Context, fileID uuid.UUID) (database.File, error) {
117119
fetches.Add(1)
118-
// Wait long enough before returning to make sure that all of the goroutines
120+
// Wait long enough before returning to make sure that all the goroutines
119121
// will be waiting in line, ensuring that no one duplicated a fetch.
120122
time.Sleep(testutil.IntervalMedium)
121-
return files.CacheEntryValue{FS: emptyFS, Size: fileSize}, nil
122-
}, reg, &coderdtest.FakeAuthorizer{})
123+
return database.File{
124+
Data: make([]byte, fileSize),
125+
}, nil
126+
}).AnyTimes()
127+
128+
c := files.New(reg, &coderdtest.FakeAuthorizer{})
123129

124130
batches := 1000
125131
groups := make([]*errgroup.Group, 0, batches)
@@ -137,7 +143,7 @@ func TestConcurrency(t *testing.T) {
137143
g.Go(func() error {
138144
// We don't bother to Release these references because the Cache will be
139145
// released at the end of the test anyway.
140-
_, err := c.Acquire(ctx, id)
146+
_, err := c.Acquire(ctx, dbM, id)
141147
return err
142148
})
143149
}
@@ -164,14 +170,15 @@ func TestRelease(t *testing.T) {
164170
ctx := dbauthz.AsFileReader(t.Context())
165171

166172
const fileSize = 10
167-
emptyFS := afero.NewIOFS(afero.NewReadOnlyFs(afero.NewMemMapFs()))
168173
reg := prometheus.NewRegistry()
169-
c := files.New(func(_ context.Context, _ uuid.UUID) (files.CacheEntryValue, error) {
170-
return files.CacheEntryValue{
171-
FS: emptyFS,
172-
Size: fileSize,
174+
dbM := dbmock.NewMockStore(gomock.NewController(t))
175+
dbM.EXPECT().GetFileByID(gomock.Any(), gomock.Any()).DoAndReturn(func(mTx context.Context, fileID uuid.UUID) (database.File, error) {
176+
return database.File{
177+
Data: make([]byte, fileSize),
173178
}, nil
174-
}, reg, &coderdtest.FakeAuthorizer{})
179+
}).AnyTimes()
180+
181+
c := files.New(reg, &coderdtest.FakeAuthorizer{})
175182

176183
batches := 100
177184
ids := make([]uuid.UUID, 0, batches)
@@ -184,9 +191,8 @@ func TestRelease(t *testing.T) {
184191
batchSize := 10
185192
for openedIdx, id := range ids {
186193
for batchIdx := range batchSize {
187-
it, err := c.Acquire(ctx, id)
194+
it, err := c.Acquire(ctx, dbM, id)
188195
require.NoError(t, err)
189-
require.Equal(t, emptyFS, it.FS)
190196
releases[id] = append(releases[id], it.Close)
191197

192198
// Each time a new file is opened, the metrics should be updated as so:
@@ -257,7 +263,7 @@ func cacheAuthzSetup(t *testing.T) (database.Store, *files.Cache, *coderdtest.Re
257263

258264
// Dbauthz wrap the db
259265
db = dbauthz.New(db, rec, logger, coderdtest.AccessControlStorePointer())
260-
c := files.NewFromStore(db, reg, rec)
266+
c := files.New(reg, rec)
261267
return db, c, rec
262268
}
263269

‎coderd/files/closer.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import (
66

77
"github.com/google/uuid"
88
"golang.org/x/xerrors"
9+
10+
"github.com/coder/coder/v2/coderd/database"
911
)
1012

1113
// CacheCloser is a cache wrapper used to close all acquired files.
@@ -38,15 +40,15 @@ func (c *CacheCloser) Close() {
3840
c.closers = nil
3941
}
4042

41-
func (c *CacheCloser) Acquire(ctx context.Context, fileID uuid.UUID) (*CloseFS, error) {
43+
func (c *CacheCloser) Acquire(ctx context.Context, db database.Store, fileID uuid.UUID) (*CloseFS, error) {
4244
c.mu.Lock()
4345
defer c.mu.Unlock()
4446

4547
if c.cache == nil {
4648
return nil, xerrors.New("cache is closed, and cannot acquire new files")
4749
}
4850

49-
f, err := c.cache.Acquire(ctx, fileID)
51+
f, err := c.cache.Acquire(ctx, db, fileID)
5052
if err != nil {
5153
return nil, err
5254
}

0 commit comments

Comments
 (0)