mirror of
https://github.com/coder/coder.git
synced 2026-09-01 14:53:15 +08:00
Zach
953091c7bc
Migrate `wg.Add(1); go func() { defer wg.Done(); ... }()` to
`wg.Go(func() { ... })` in tests.
Where the prior pattern passed the loop variable explicitly via a
closure parameter (`go func(id int) { ... }(i)`), drop the parameter and
reference the loop variable directly. Per-iteration loop variables since
Go 1.22 make this safe.
593 lines
18 KiB
Go
593 lines
18 KiB
Go
package boundaryusage_test
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"testing"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/goleak"
|
|
|
|
"github.com/coder/coder/v2/coderd/boundaryusage"
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
|
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
|
"github.com/coder/coder/v2/testutil"
|
|
)
|
|
|
|
func TestMain(m *testing.M) {
|
|
goleak.VerifyTestMain(m, testutil.GoleakOptions...)
|
|
}
|
|
|
|
func TestTracker_New(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
require.NotNil(t, tracker)
|
|
}
|
|
|
|
func TestTracker_Track_Single(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
workspaceID := uuid.New()
|
|
ownerID := uuid.New()
|
|
replicaID := uuid.New()
|
|
|
|
tracker.Track(workspaceID, ownerID, 5, 2)
|
|
|
|
err := tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
// Verify the data was written correctly.
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(1), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(1), summary.UniqueUsers)
|
|
require.Equal(t, int64(5), summary.AllowedRequests)
|
|
require.Equal(t, int64(2), summary.DeniedRequests)
|
|
}
|
|
|
|
func TestTracker_Track_DuplicateWorkspaceUser(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
workspaceID := uuid.New()
|
|
ownerID := uuid.New()
|
|
replicaID := uuid.New()
|
|
|
|
// Track same workspace/user multiple times.
|
|
tracker.Track(workspaceID, ownerID, 3, 1)
|
|
tracker.Track(workspaceID, ownerID, 4, 2)
|
|
tracker.Track(workspaceID, ownerID, 2, 0)
|
|
|
|
err := tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(1), summary.UniqueWorkspaces, "should be 1 unique workspace")
|
|
require.Equal(t, int64(1), summary.UniqueUsers, "should be 1 unique user")
|
|
require.Equal(t, int64(9), summary.AllowedRequests, "should accumulate: 3+4+2=9")
|
|
require.Equal(t, int64(3), summary.DeniedRequests, "should accumulate: 1+2+0=3")
|
|
}
|
|
|
|
func TestTracker_Track_MultipleWorkspacesUsers(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
replicaID := uuid.New()
|
|
|
|
// Track 3 different workspaces with 2 different users.
|
|
workspace1, workspace2, workspace3 := uuid.New(), uuid.New(), uuid.New()
|
|
user1, user2 := uuid.New(), uuid.New()
|
|
|
|
tracker.Track(workspace1, user1, 1, 0)
|
|
tracker.Track(workspace2, user1, 2, 1)
|
|
tracker.Track(workspace3, user2, 3, 2)
|
|
|
|
err := tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(3), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(2), summary.UniqueUsers)
|
|
require.Equal(t, int64(6), summary.AllowedRequests)
|
|
require.Equal(t, int64(3), summary.DeniedRequests)
|
|
}
|
|
|
|
func TestTracker_Track_Concurrent(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
replicaID := uuid.New()
|
|
|
|
const numGoroutines = 100
|
|
const requestsPerGoroutine = 10
|
|
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < numGoroutines; i++ {
|
|
wg.Go(func() {
|
|
workspaceID := uuid.New()
|
|
ownerID := uuid.New()
|
|
for j := 0; j < requestsPerGoroutine; j++ {
|
|
tracker.Track(workspaceID, ownerID, 1, 1)
|
|
}
|
|
})
|
|
}
|
|
wg.Wait()
|
|
|
|
err := tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(numGoroutines), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(numGoroutines), summary.UniqueUsers)
|
|
require.Equal(t, int64(numGoroutines*requestsPerGoroutine), summary.AllowedRequests)
|
|
require.Equal(t, int64(numGoroutines*requestsPerGoroutine), summary.DeniedRequests)
|
|
}
|
|
|
|
func TestTracker_FlushToDB_Accumulates(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
replicaID := uuid.New()
|
|
workspaceID := uuid.New()
|
|
ownerID := uuid.New()
|
|
|
|
// First flush is an insert, resets unique counts (new period).
|
|
tracker.Track(workspaceID, ownerID, 5, 3)
|
|
err := tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
// Track & flush more data. Same workspace/user, so unique counts stay at 1.
|
|
tracker.Track(workspaceID, ownerID, 2, 1)
|
|
err = tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
// Track & flush even more data to continue accumulation.
|
|
tracker.Track(workspaceID, ownerID, 3, 2)
|
|
err = tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(1), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(1), summary.UniqueUsers)
|
|
require.Equal(t, int64(5+2+3), summary.AllowedRequests)
|
|
require.Equal(t, int64(3+1+2), summary.DeniedRequests)
|
|
}
|
|
|
|
func TestTracker_FlushToDB_NewPeriod(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
replicaID := uuid.New()
|
|
workspaceID := uuid.New()
|
|
ownerID := uuid.New()
|
|
|
|
tracker.Track(workspaceID, ownerID, 10, 5)
|
|
|
|
// First flush.
|
|
err := tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
// Simulate telemetry reset (new period).
|
|
_, err = db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
|
|
// Track new data.
|
|
workspace2 := uuid.New()
|
|
owner2 := uuid.New()
|
|
tracker.Track(workspace2, owner2, 3, 1)
|
|
|
|
// Flushing again should detect new period and reset in-memory stats.
|
|
err = tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
// The summary should only contain the new data after reset.
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(1), summary.UniqueWorkspaces, "should only count new workspace")
|
|
require.Equal(t, int64(1), summary.UniqueUsers, "should only count new user")
|
|
require.Equal(t, int64(3), summary.AllowedRequests, "should only count new requests")
|
|
require.Equal(t, int64(1), summary.DeniedRequests, "should only count new requests")
|
|
}
|
|
|
|
func TestTracker_FlushToDB_NoActivity(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
replicaID := uuid.New()
|
|
|
|
err := tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
// Verify nothing was written to DB.
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(0), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(0), summary.AllowedRequests)
|
|
}
|
|
|
|
func TestUpsertBoundaryUsageStats_Insert(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := dbauthz.AsBoundaryUsageTracker(context.Background())
|
|
|
|
replicaID := uuid.New()
|
|
|
|
// Set different values for delta vs cumulative to verify INSERT uses delta.
|
|
newPeriod, err := db.UpsertBoundaryUsageStats(ctx, database.UpsertBoundaryUsageStatsParams{
|
|
ReplicaID: replicaID,
|
|
UniqueWorkspacesDelta: 5,
|
|
UniqueUsersDelta: 3,
|
|
UniqueWorkspacesCount: 999, // should be ignored on INSERT
|
|
UniqueUsersCount: 999, // should be ignored on INSERT
|
|
AllowedRequests: 100,
|
|
DeniedRequests: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
require.True(t, newPeriod, "should return true for insert")
|
|
|
|
// Verify INSERT used the delta values, not cumulative.
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(ctx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(5), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(3), summary.UniqueUsers)
|
|
}
|
|
|
|
func TestUpsertBoundaryUsageStats_Update(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := dbauthz.AsBoundaryUsageTracker(context.Background())
|
|
|
|
replicaID := uuid.New()
|
|
|
|
// First insert uses delta fields.
|
|
_, err := db.UpsertBoundaryUsageStats(ctx, database.UpsertBoundaryUsageStatsParams{
|
|
ReplicaID: replicaID,
|
|
UniqueWorkspacesDelta: 5,
|
|
UniqueUsersDelta: 3,
|
|
AllowedRequests: 100,
|
|
DeniedRequests: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Second upsert (update). Set different delta vs cumulative to verify UPDATE uses cumulative.
|
|
newPeriod, err := db.UpsertBoundaryUsageStats(ctx, database.UpsertBoundaryUsageStatsParams{
|
|
ReplicaID: replicaID,
|
|
UniqueWorkspacesCount: 8, // cumulative, should be used
|
|
UniqueUsersCount: 5, // cumulative, should be used
|
|
AllowedRequests: 200,
|
|
DeniedRequests: 20,
|
|
})
|
|
require.NoError(t, err)
|
|
require.False(t, newPeriod, "should return false for update")
|
|
|
|
// Verify UPDATE used cumulative values.
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(ctx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(8), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(5), summary.UniqueUsers)
|
|
require.Equal(t, int64(100+200), summary.AllowedRequests)
|
|
require.Equal(t, int64(10+20), summary.DeniedRequests)
|
|
}
|
|
|
|
func TestGetAndResetBoundaryUsageSummary_MultipleReplicas(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := dbauthz.AsBoundaryUsageTracker(context.Background())
|
|
|
|
replica1 := uuid.New()
|
|
replica2 := uuid.New()
|
|
replica3 := uuid.New()
|
|
|
|
// Insert stats for 3 replicas. Delta fields are used for INSERT.
|
|
_, err := db.UpsertBoundaryUsageStats(ctx, database.UpsertBoundaryUsageStatsParams{
|
|
ReplicaID: replica1,
|
|
UniqueWorkspacesDelta: 10,
|
|
UniqueUsersDelta: 5,
|
|
AllowedRequests: 100,
|
|
DeniedRequests: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
_, err = db.UpsertBoundaryUsageStats(ctx, database.UpsertBoundaryUsageStatsParams{
|
|
ReplicaID: replica2,
|
|
UniqueWorkspacesDelta: 15,
|
|
UniqueUsersDelta: 8,
|
|
AllowedRequests: 150,
|
|
DeniedRequests: 15,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
_, err = db.UpsertBoundaryUsageStats(ctx, database.UpsertBoundaryUsageStatsParams{
|
|
ReplicaID: replica3,
|
|
UniqueWorkspacesDelta: 20,
|
|
UniqueUsersDelta: 12,
|
|
AllowedRequests: 200,
|
|
DeniedRequests: 20,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(ctx, 60000)
|
|
require.NoError(t, err)
|
|
|
|
// Verify aggregation (SUM of all replicas).
|
|
require.Equal(t, int64(45), summary.UniqueWorkspaces) // 10 + 15 + 20
|
|
require.Equal(t, int64(25), summary.UniqueUsers) // 5 + 8 + 12
|
|
require.Equal(t, int64(450), summary.AllowedRequests) // 100 + 150 + 200
|
|
require.Equal(t, int64(45), summary.DeniedRequests) // 10 + 15 + 20
|
|
}
|
|
|
|
func TestGetAndResetBoundaryUsageSummary_Empty(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := dbauthz.AsBoundaryUsageTracker(context.Background())
|
|
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(ctx, 60000)
|
|
require.NoError(t, err)
|
|
|
|
// COALESCE should return 0 for all columns.
|
|
require.Equal(t, int64(0), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(0), summary.UniqueUsers)
|
|
require.Equal(t, int64(0), summary.AllowedRequests)
|
|
require.Equal(t, int64(0), summary.DeniedRequests)
|
|
}
|
|
|
|
func TestGetAndResetBoundaryUsageSummary_DeletesData(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := dbauthz.AsBoundaryUsageTracker(context.Background())
|
|
|
|
// Insert stats for multiple replicas. Delta fields are used for INSERT.
|
|
for i := 0; i < 5; i++ {
|
|
_, err := db.UpsertBoundaryUsageStats(ctx, database.UpsertBoundaryUsageStatsParams{
|
|
ReplicaID: uuid.New(),
|
|
UniqueWorkspacesDelta: int64(i + 1),
|
|
UniqueUsersDelta: int64(i + 1),
|
|
AllowedRequests: int64((i + 1) * 10),
|
|
DeniedRequests: int64(i + 1),
|
|
})
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// Should return the summary AND delete all data.
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(ctx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(1+2+3+4+5), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(10+20+30+40+50), summary.AllowedRequests)
|
|
|
|
// Verify all data is gone.
|
|
summary, err = db.GetAndResetBoundaryUsageSummary(ctx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(0), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(0), summary.AllowedRequests)
|
|
}
|
|
|
|
func TestTracker_TelemetryCycle(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
|
|
// Simulate 3 replicas.
|
|
tracker1 := boundaryusage.NewTracker()
|
|
tracker2 := boundaryusage.NewTracker()
|
|
tracker3 := boundaryusage.NewTracker()
|
|
|
|
replica1 := uuid.New()
|
|
replica2 := uuid.New()
|
|
replica3 := uuid.New()
|
|
|
|
// Each tracker records different workspaces/users.
|
|
tracker1.Track(uuid.New(), uuid.New(), 10, 1)
|
|
tracker1.Track(uuid.New(), uuid.New(), 15, 2)
|
|
|
|
tracker2.Track(uuid.New(), uuid.New(), 20, 3)
|
|
tracker2.Track(uuid.New(), uuid.New(), 25, 4)
|
|
tracker2.Track(uuid.New(), uuid.New(), 30, 5)
|
|
|
|
tracker3.Track(uuid.New(), uuid.New(), 5, 0)
|
|
|
|
// All replicas flush to database.
|
|
require.NoError(t, tracker1.FlushToDB(ctx, db, replica1))
|
|
require.NoError(t, tracker2.FlushToDB(ctx, db, replica2))
|
|
require.NoError(t, tracker3.FlushToDB(ctx, db, replica3))
|
|
|
|
// Telemetry aggregates and resets (simulating telemetry report sent).
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
|
|
// Verify aggregation.
|
|
require.Equal(t, int64(6), summary.UniqueWorkspaces) // 2 + 3 + 1
|
|
require.Equal(t, int64(6), summary.UniqueUsers) // 2 + 3 + 1
|
|
require.Equal(t, int64(105), summary.AllowedRequests) // 25 + 75 + 5
|
|
require.Equal(t, int64(15), summary.DeniedRequests) // 3 + 12 + 0
|
|
|
|
// Next flush from trackers should detect new period.
|
|
tracker1.Track(uuid.New(), uuid.New(), 1, 0)
|
|
require.NoError(t, tracker1.FlushToDB(ctx, db, replica1))
|
|
|
|
// Verify trackers reset their in-memory state.
|
|
summary, err = db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(1), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(1), summary.AllowedRequests)
|
|
}
|
|
|
|
func TestTracker_FlushToDB_NoStaleDataAfterReset(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
replicaID := uuid.New()
|
|
workspaceID := uuid.New()
|
|
ownerID := uuid.New()
|
|
|
|
// Track some data and flush.
|
|
tracker.Track(workspaceID, ownerID, 10, 5)
|
|
err := tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
// Simulate telemetry reset (new period) - this also verifies the data.
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(1), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(10), summary.AllowedRequests)
|
|
|
|
// Flush again without any new Track() calls. This should not write stale
|
|
// data back to the DB.
|
|
err = tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
// Summary should be empty (no stale data written).
|
|
summary, err = db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(0), summary.UniqueWorkspaces)
|
|
require.Equal(t, int64(0), summary.UniqueUsers)
|
|
require.Equal(t, int64(0), summary.AllowedRequests)
|
|
require.Equal(t, int64(0), summary.DeniedRequests)
|
|
}
|
|
|
|
func TestTracker_ConcurrentFlushAndTrack(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
replicaID := uuid.New()
|
|
|
|
const numOperations = 50
|
|
|
|
var wg sync.WaitGroup
|
|
|
|
// Goroutine 1: Continuously track.
|
|
wg.Go(func() {
|
|
for i := 0; i < numOperations; i++ {
|
|
tracker.Track(uuid.New(), uuid.New(), 1, 1)
|
|
}
|
|
})
|
|
|
|
// Goroutine 2: Continuously flush.
|
|
wg.Go(func() {
|
|
for i := 0; i < numOperations; i++ {
|
|
_ = tracker.FlushToDB(ctx, db, replicaID)
|
|
}
|
|
})
|
|
|
|
wg.Wait()
|
|
|
|
// Final flush to capture any remaining data.
|
|
require.NoError(t, tracker.FlushToDB(ctx, db, replicaID))
|
|
|
|
// Verify stats are non-negative.
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.GreaterOrEqual(t, summary.AllowedRequests, int64(0))
|
|
require.GreaterOrEqual(t, summary.DeniedRequests, int64(0))
|
|
}
|
|
|
|
// trackDuringUpsertDB wraps a database.Store to call Track() during the
|
|
// UpsertBoundaryUsageStats operation, simulating a concurrent Track() call.
|
|
type trackDuringUpsertDB struct {
|
|
database.Store
|
|
tracker *boundaryusage.Tracker
|
|
workspaceID uuid.UUID
|
|
userID uuid.UUID
|
|
}
|
|
|
|
func (s *trackDuringUpsertDB) InTx(fn func(database.Store) error, opts *database.TxOptions) error {
|
|
return s.Store.InTx(func(tx database.Store) error {
|
|
return fn(&trackDuringUpsertDB{
|
|
Store: tx,
|
|
tracker: s.tracker,
|
|
workspaceID: s.workspaceID,
|
|
userID: s.userID,
|
|
})
|
|
}, opts)
|
|
}
|
|
|
|
func (s *trackDuringUpsertDB) UpsertBoundaryUsageStats(ctx context.Context, arg database.UpsertBoundaryUsageStatsParams) (bool, error) {
|
|
s.tracker.Track(s.workspaceID, s.userID, 20, 10)
|
|
return s.Store.UpsertBoundaryUsageStats(ctx, arg)
|
|
}
|
|
|
|
func TestTracker_TrackDuringFlush(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, _ := dbtestutil.NewDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
boundaryCtx := dbauthz.AsBoundaryUsageTracker(ctx)
|
|
|
|
tracker := boundaryusage.NewTracker()
|
|
replicaID := uuid.New()
|
|
|
|
// Track some initial data.
|
|
tracker.Track(uuid.New(), uuid.New(), 10, 5)
|
|
|
|
trackingDB := &trackDuringUpsertDB{
|
|
Store: db,
|
|
tracker: tracker,
|
|
workspaceID: uuid.New(),
|
|
userID: uuid.New(),
|
|
}
|
|
|
|
// Flush will call Track() during the DB operation.
|
|
err := tracker.FlushToDB(ctx, trackingDB, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
// Second flush captures the Track() that happened during the first flush.
|
|
err = tracker.FlushToDB(ctx, db, replicaID)
|
|
require.NoError(t, err)
|
|
|
|
// Verify both flushes are in the summary.
|
|
summary, err := db.GetAndResetBoundaryUsageSummary(boundaryCtx, 60000)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(10+20), summary.AllowedRequests)
|
|
require.Equal(t, int64(5+10), summary.DeniedRequests)
|
|
}
|