mirror of
https://github.com/coder/coder.git
synced 2026-06-03 13:08:25 +00:00
Spike Curtis
bddb808b25
Fixes all our Go file imports to match the preferred spec that we've _mostly_ been using. For example: ``` import ( "context" "time" "github.com/prometheus/client_golang/prometheus" "golang.org/x/xerrors" "gopkg.in/natefinch/lumberjack.v2" "cdr.dev/slog/v3" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/serpent" ) ``` 3 groups: standard library, 3rd partly libs, Coder libs. This PR makes the change across the codebase. The PR in the stack above modifies our formatting to maintain this state of affairs, and is a separate PR so it's possible to review that one in detail.
365 lines
12 KiB
Go
365 lines
12 KiB
Go
package files_test
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/mock/gomock"
|
|
"golang.org/x/sync/errgroup"
|
|
|
|
"cdr.dev/slog/v3/sloggers/slogtest"
|
|
"github.com/coder/coder/v2/coderd/coderdtest"
|
|
"github.com/coder/coder/v2/coderd/coderdtest/promhelp"
|
|
"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/dbmock"
|
|
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
|
"github.com/coder/coder/v2/coderd/files"
|
|
"github.com/coder/coder/v2/coderd/rbac"
|
|
"github.com/coder/coder/v2/coderd/rbac/policy"
|
|
"github.com/coder/coder/v2/testutil"
|
|
)
|
|
|
|
func TestCancelledFetch(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
fileID := uuid.New()
|
|
dbM := dbmock.NewMockStore(gomock.NewController(t))
|
|
|
|
// The file fetch should succeed.
|
|
dbM.EXPECT().GetFileByID(gomock.Any(), gomock.Any()).DoAndReturn(func(mTx context.Context, fileID uuid.UUID) (database.File, error) {
|
|
return database.File{
|
|
ID: fileID,
|
|
Data: make([]byte, 100),
|
|
}, nil
|
|
})
|
|
|
|
cache := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
|
|
|
|
// Cancel the context for the first call; should fail.
|
|
ctx, cancel := context.WithCancel(dbauthz.AsFileReader(testutil.Context(t, testutil.WaitShort)))
|
|
cancel()
|
|
_, err := cache.Acquire(ctx, dbM, fileID)
|
|
assert.ErrorIs(t, err, context.Canceled)
|
|
}
|
|
|
|
// TestCancelledConcurrentFetch runs 2 Acquire calls. The first has a canceled
|
|
// context and will get a ctx.Canceled error. The second call should get a warmfirst error and try to fetch the file
|
|
// again, which should succeed.
|
|
func TestCancelledConcurrentFetch(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
fileID := uuid.New()
|
|
dbM := dbmock.NewMockStore(gomock.NewController(t))
|
|
|
|
// The file fetch should succeed.
|
|
dbM.EXPECT().GetFileByID(gomock.Any(), gomock.Any()).DoAndReturn(func(mTx context.Context, fileID uuid.UUID) (database.File, error) {
|
|
return database.File{
|
|
ID: fileID,
|
|
Data: make([]byte, 100),
|
|
}, nil
|
|
})
|
|
|
|
cache := files.LeakCache{Cache: files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})}
|
|
|
|
ctx := dbauthz.AsFileReader(testutil.Context(t, testutil.WaitShort))
|
|
|
|
// Cancel the context for the first call; should fail.
|
|
canceledCtx, cancel := context.WithCancel(ctx)
|
|
cancel()
|
|
_, err := cache.Acquire(canceledCtx, dbM, fileID)
|
|
require.ErrorIs(t, err, context.Canceled)
|
|
|
|
// Second call, that should succeed without fetching from the database again
|
|
// since the cache should be populated by the fetch the first request started
|
|
// even if it doesn't wait for completion.
|
|
_, err = cache.Acquire(ctx, dbM, fileID)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
func TestConcurrentFetch(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
fileID := uuid.New()
|
|
|
|
// Only allow one call, which should succeed
|
|
dbM := dbmock.NewMockStore(gomock.NewController(t))
|
|
dbM.EXPECT().GetFileByID(gomock.Any(), gomock.Any()).DoAndReturn(func(mTx context.Context, fileID uuid.UUID) (database.File, error) {
|
|
return database.File{ID: fileID}, nil
|
|
})
|
|
|
|
cache := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
|
|
ctx := dbauthz.AsFileReader(testutil.Context(t, testutil.WaitShort))
|
|
|
|
// Expect 2 calls to Acquire before we continue the test
|
|
var wg sync.WaitGroup
|
|
|
|
wg.Add(2)
|
|
for range 2 {
|
|
// TODO: wg.Go in Go 1.25
|
|
go func() {
|
|
defer wg.Done()
|
|
_, err := cache.Acquire(ctx, dbM, fileID)
|
|
assert.NoError(t, err)
|
|
}()
|
|
}
|
|
|
|
// Wait for both go routines to assert their errors and finish.
|
|
wg.Wait()
|
|
require.Equal(t, 1, cache.Count())
|
|
}
|
|
|
|
// nolint:paralleltest,tparallel // Serially testing is easier
|
|
func TestCacheRBAC(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, cache, rec := cacheAuthzSetup(t)
|
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
|
|
|
file := dbgen.File(t, db, database.File{})
|
|
|
|
nobodyID := uuid.New()
|
|
nobody := dbauthz.As(ctx, rbac.Subject{
|
|
ID: nobodyID.String(),
|
|
Roles: rbac.Roles{},
|
|
Scope: rbac.ScopeAll,
|
|
})
|
|
|
|
userID := uuid.New()
|
|
userReader := dbauthz.As(ctx, rbac.Subject{
|
|
ID: userID.String(),
|
|
Roles: rbac.Roles{
|
|
must(rbac.RoleByName(rbac.RoleTemplateAdmin())),
|
|
},
|
|
Scope: rbac.ScopeAll,
|
|
})
|
|
|
|
cacheReader := dbauthz.AsFileReader(ctx)
|
|
|
|
t.Run("NoRolesOpen", func(t *testing.T) {
|
|
// Ensure start is clean
|
|
require.Equal(t, 0, cache.Count())
|
|
rec.Reset()
|
|
|
|
_, err := cache.Acquire(nobody, db, file.ID)
|
|
require.Error(t, err)
|
|
require.True(t, rbac.IsUnauthorizedError(err))
|
|
|
|
// Ensure that the cache is empty
|
|
require.Equal(t, 0, cache.Count())
|
|
|
|
// Check the assertions
|
|
rec.AssertActorID(t, nobodyID.String(), rec.Pair(policy.ActionRead, file))
|
|
rec.AssertActorID(t, rbac.SubjectTypeFileReaderID, rec.Pair(policy.ActionRead, file))
|
|
})
|
|
|
|
t.Run("CacheHasFile", func(t *testing.T) {
|
|
rec.Reset()
|
|
require.Equal(t, 0, cache.Count())
|
|
|
|
// Read the file with a file reader to put it into the cache.
|
|
a, err := cache.Acquire(cacheReader, db, file.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, cache.Count())
|
|
|
|
// "nobody" should not be able to read the file.
|
|
_, err = cache.Acquire(nobody, db, file.ID)
|
|
require.Error(t, err)
|
|
require.True(t, rbac.IsUnauthorizedError(err))
|
|
require.Equal(t, 1, cache.Count())
|
|
|
|
// UserReader can
|
|
b, err := cache.Acquire(userReader, db, file.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, cache.Count())
|
|
|
|
a.Close()
|
|
b.Close()
|
|
require.Equal(t, 0, cache.Count())
|
|
|
|
rec.AssertActorID(t, nobodyID.String(), rec.Pair(policy.ActionRead, file))
|
|
rec.AssertActorID(t, rbac.SubjectTypeFileReaderID, rec.Pair(policy.ActionRead, file))
|
|
rec.AssertActorID(t, userID.String(), rec.Pair(policy.ActionRead, file))
|
|
})
|
|
}
|
|
|
|
func cachePromMetricName(metric string) string {
|
|
return "coderd_file_cache_" + metric
|
|
}
|
|
|
|
func TestConcurrency(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := dbauthz.AsFileReader(t.Context())
|
|
|
|
const fileSize = 10
|
|
var fetches atomic.Int64
|
|
reg := prometheus.NewRegistry()
|
|
|
|
dbM := dbmock.NewMockStore(gomock.NewController(t))
|
|
dbM.EXPECT().GetFileByID(gomock.Any(), gomock.Any()).DoAndReturn(func(mTx context.Context, fileID uuid.UUID) (database.File, error) {
|
|
fetches.Add(1)
|
|
// Wait long enough before returning to make sure that all the goroutines
|
|
// will be waiting in line, ensuring that no one duplicated a fetch.
|
|
time.Sleep(testutil.IntervalMedium)
|
|
return database.File{
|
|
Data: make([]byte, fileSize),
|
|
}, nil
|
|
}).AnyTimes()
|
|
|
|
c := files.New(reg, &coderdtest.FakeAuthorizer{})
|
|
|
|
batches := 1000
|
|
groups := make([]*errgroup.Group, 0, batches)
|
|
for range batches {
|
|
groups = append(groups, new(errgroup.Group))
|
|
}
|
|
|
|
// Call Acquire with a unique ID per batch, many times per batch, with many
|
|
// batches all in parallel. This is pretty much the worst-case scenario:
|
|
// thousands of concurrent reads, with both warm and cold loads happening.
|
|
batchSize := 10
|
|
for _, g := range groups {
|
|
id := uuid.New()
|
|
for range batchSize {
|
|
g.Go(func() error {
|
|
// We don't bother to Release these references because the Cache will be
|
|
// released at the end of the test anyway.
|
|
_, err := c.Acquire(ctx, dbM, id)
|
|
return err
|
|
})
|
|
}
|
|
}
|
|
|
|
for _, g := range groups {
|
|
require.NoError(t, g.Wait())
|
|
}
|
|
require.Equal(t, int64(batches), fetches.Load())
|
|
|
|
// Verify all the counts & metrics are correct.
|
|
require.Equal(t, batches, c.Count())
|
|
require.Equal(t, batches*fileSize, promhelp.GaugeValue(t, reg, cachePromMetricName("open_files_size_bytes_current"), nil))
|
|
require.Equal(t, batches*fileSize, promhelp.CounterValue(t, reg, cachePromMetricName("open_files_size_bytes_total"), nil))
|
|
require.Equal(t, batches, promhelp.GaugeValue(t, reg, cachePromMetricName("open_files_current"), nil))
|
|
require.Equal(t, batches, promhelp.CounterValue(t, reg, cachePromMetricName("open_files_total"), nil))
|
|
require.Equal(t, batches*batchSize, promhelp.GaugeValue(t, reg, cachePromMetricName("open_file_refs_current"), nil))
|
|
hit, miss := promhelp.CounterValue(t, reg, cachePromMetricName("open_file_refs_total"), prometheus.Labels{"hit": "false"}),
|
|
promhelp.CounterValue(t, reg, cachePromMetricName("open_file_refs_total"), prometheus.Labels{"hit": "true"})
|
|
require.Equal(t, batches*batchSize, hit+miss)
|
|
}
|
|
|
|
func TestRelease(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := dbauthz.AsFileReader(t.Context())
|
|
|
|
const fileSize = 10
|
|
reg := prometheus.NewRegistry()
|
|
dbM := dbmock.NewMockStore(gomock.NewController(t))
|
|
dbM.EXPECT().GetFileByID(gomock.Any(), gomock.Any()).DoAndReturn(func(mTx context.Context, fileID uuid.UUID) (database.File, error) {
|
|
return database.File{
|
|
Data: make([]byte, fileSize),
|
|
}, nil
|
|
}).AnyTimes()
|
|
|
|
c := files.New(reg, &coderdtest.FakeAuthorizer{})
|
|
|
|
batches := 100
|
|
ids := make([]uuid.UUID, 0, batches)
|
|
for range batches {
|
|
ids = append(ids, uuid.New())
|
|
}
|
|
|
|
releases := make(map[uuid.UUID][]func(), 0)
|
|
// Acquire a bunch of references
|
|
batchSize := 10
|
|
for openedIdx, id := range ids {
|
|
for batchIdx := range batchSize {
|
|
it, err := c.Acquire(ctx, dbM, id)
|
|
require.NoError(t, err)
|
|
releases[id] = append(releases[id], it.Close)
|
|
|
|
// Each time a new file is opened, the metrics should be updated as so:
|
|
opened := openedIdx + 1
|
|
// Number of unique files opened is equal to the idx of the ids.
|
|
require.Equal(t, opened, c.Count())
|
|
require.Equal(t, opened, promhelp.GaugeValue(t, reg, cachePromMetricName("open_files_current"), nil))
|
|
// Current file size is unique files * file size.
|
|
require.Equal(t, opened*fileSize, promhelp.GaugeValue(t, reg, cachePromMetricName("open_files_size_bytes_current"), nil))
|
|
// The number of refs is the current iteration of both loops.
|
|
require.Equal(t, ((opened-1)*batchSize)+(batchIdx+1), promhelp.GaugeValue(t, reg, cachePromMetricName("open_file_refs_current"), nil))
|
|
}
|
|
}
|
|
|
|
// Make sure cache is fully loaded
|
|
require.Equal(t, c.Count(), batches)
|
|
|
|
// Now release all of the references
|
|
for closedIdx, id := range ids {
|
|
stillOpen := len(ids) - closedIdx
|
|
for closingIdx := range batchSize {
|
|
releases[id][0]()
|
|
releases[id] = releases[id][1:]
|
|
|
|
// Each time a file is released, the metrics should decrement the file refs
|
|
require.Equal(t, (stillOpen*batchSize)-(closingIdx+1), promhelp.GaugeValue(t, reg, cachePromMetricName("open_file_refs_current"), nil))
|
|
|
|
closed := closingIdx+1 == batchSize
|
|
if closed {
|
|
continue
|
|
}
|
|
|
|
// File ref still exists, so the counts should not change yet.
|
|
require.Equal(t, stillOpen, c.Count())
|
|
require.Equal(t, stillOpen, promhelp.GaugeValue(t, reg, cachePromMetricName("open_files_current"), nil))
|
|
require.Equal(t, stillOpen*fileSize, promhelp.GaugeValue(t, reg, cachePromMetricName("open_files_size_bytes_current"), nil))
|
|
}
|
|
}
|
|
|
|
// ...and make sure that the cache has emptied itself.
|
|
require.Equal(t, c.Count(), 0)
|
|
|
|
// Verify all the counts & metrics are correct.
|
|
// All existing files are closed
|
|
require.Equal(t, 0, c.Count())
|
|
require.Equal(t, 0, promhelp.GaugeValue(t, reg, cachePromMetricName("open_files_size_bytes_current"), nil))
|
|
require.Equal(t, 0, promhelp.GaugeValue(t, reg, cachePromMetricName("open_files_current"), nil))
|
|
require.Equal(t, 0, promhelp.GaugeValue(t, reg, cachePromMetricName("open_file_refs_current"), nil))
|
|
|
|
// Total counts remain
|
|
require.Equal(t, batches*fileSize, promhelp.CounterValue(t, reg, cachePromMetricName("open_files_size_bytes_total"), nil))
|
|
require.Equal(t, batches, promhelp.CounterValue(t, reg, cachePromMetricName("open_files_total"), nil))
|
|
}
|
|
|
|
func cacheAuthzSetup(t *testing.T) (database.Store, *files.Cache, *coderdtest.RecordingAuthorizer) {
|
|
t.Helper()
|
|
|
|
logger := slogtest.Make(t, &slogtest.Options{})
|
|
reg := prometheus.NewRegistry()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
authz := rbac.NewAuthorizer(reg)
|
|
rec := &coderdtest.RecordingAuthorizer{
|
|
Called: nil,
|
|
Wrapped: authz,
|
|
}
|
|
|
|
// Dbauthz wrap the db
|
|
db = dbauthz.New(db, rec, logger, coderdtest.AccessControlStorePointer())
|
|
c := files.New(reg, rec)
|
|
return db, c, rec
|
|
}
|
|
|
|
func must[T any](t T, err error) T {
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return t
|
|
}
|