coder · f0ssel · May 29, 2024 · May 29, 2024 · May 29, 2024 · May 29, 2024
diff --git a/cli/portforward_test.go b/cli/portforward_test.go
@@ -21,7 +21,6 @@ import (
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
-	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/pty/ptytest"
 	"github.com/coder/coder/v2/testutil"
@@ -97,12 +96,7 @@ func TestPortForward(t *testing.T) {
 	// Setup agent once to be shared between test-cases (avoid expensive
 	// non-parallel setup).
 	var (
-		wuTick     = make(chan time.Time)
-		wuFlush    = make(chan int, 1)
-		client, db = coderdtest.NewWithDatabase(t, &coderdtest.Options{
-			WorkspaceUsageTrackerTick:  wuTick,
-			WorkspaceUsageTrackerFlush: wuFlush,
-		})
+		client, db         = coderdtest.NewWithDatabase(t, nil)
 		admin              = coderdtest.CreateFirstUser(t, client)
 		member, memberUser = coderdtest.CreateAnotherUser(t, client, admin.OrganizationID)
 		workspace          = runAgent(t, client, memberUser.ID, db)
@@ -155,9 +149,6 @@ func TestPortForward(t *testing.T) {
 			err = <-errC
 			require.ErrorIs(t, err, context.Canceled)
 
-			flushCtx := testutil.Context(t, testutil.WaitShort)
-			testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
-			_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
 			updated, err := client.Workspace(context.Background(), workspace.ID)
 			require.NoError(t, err)
 			require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
@@ -210,9 +201,6 @@ func TestPortForward(t *testing.T) {
 			err = <-errC
 			require.ErrorIs(t, err, context.Canceled)
 
-			flushCtx := testutil.Context(t, testutil.WaitShort)
-			testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
-			_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
 			updated, err := client.Workspace(context.Background(), workspace.ID)
 			require.NoError(t, err)
 			require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
@@ -278,9 +266,6 @@ func TestPortForward(t *testing.T) {
 		err := <-errC
 		require.ErrorIs(t, err, context.Canceled)
 
-		flushCtx := testutil.Context(t, testutil.WaitShort)
-		testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
-		_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
 		updated, err := client.Workspace(context.Background(), workspace.ID)
 		require.NoError(t, err)
 		require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)

diff --git a/cli/server.go b/cli/server.go
@@ -62,7 +62,6 @@ import (
 	"github.com/coder/coder/v2/cli/config"
 	"github.com/coder/coder/v2/coderd"
 	"github.com/coder/coder/v2/coderd/autobuild"
-	"github.com/coder/coder/v2/coderd/batchstats"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/awsiamrds"
 	"github.com/coder/coder/v2/coderd/database/dbmem"
@@ -87,7 +86,6 @@ import (
 	stringutil "github.com/coder/coder/v2/coderd/util/strings"
 	"github.com/coder/coder/v2/coderd/workspaceapps"
 	"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
-	"github.com/coder/coder/v2/coderd/workspaceusage"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/drpc"
 	"github.com/coder/coder/v2/cryptorand"
@@ -869,16 +867,6 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				options.SwaggerEndpoint = vals.Swagger.Enable.Value()
 			}
 
-			batcher, closeBatcher, err := batchstats.New(ctx,
-				batchstats.WithLogger(options.Logger.Named("batchstats")),
-				batchstats.WithStore(options.Database),
-			)
-			if err != nil {
-				return xerrors.Errorf("failed to create agent stats batcher: %w", err)
-			}
-			options.StatsBatcher = batcher
-			defer closeBatcher()
-
 			// We use a separate coderAPICloser so the Enterprise API
 			// can have its own close functions. This is cleaner
 			// than abstracting the Coder API itself.
@@ -975,13 +963,6 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 			purger := dbpurge.New(ctx, logger.Named("dbpurge"), options.Database)
 			defer purger.Close()
 
-			// Updates workspace usage
-			tracker := workspaceusage.New(options.Database,
-				workspaceusage.WithLogger(logger.Named("workspace_usage_tracker")),
-			)
-			options.WorkspaceUsageTracker = tracker
-			defer tracker.Close()
-
 			// Wrap the server in middleware that redirects to the access URL if
 			// the request is not to a local IP.
 			var ha
ED48
ndler http.Handler = coderAPI.RootHandler

diff --git a/coderd/agentapi/stats.go b/coderd/agentapi/stats.go
@@ -7,19 +7,13 @@ import (
 	"golang.org/x/xerrors"
 	"google.golang.org/protobuf/types/known/durationpb"
 
-	"github.com/google/uuid"
-
 	"cdr.dev/slog"
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/workspacestats"
 )
 
-type StatsBatcher interface {
-	Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error
-}
-
 type StatsAPI struct {
 	AgentFn                   func(context.Context) (database.WorkspaceAgent, error)
 	Database                  database.Store

diff --git a/coderd/agentapi/stats_test.go b/coderd/agentapi/stats_test.go
@@ -3,7 +3,6 @@ package agentapi_test
 import (
 	"context"
 	"database/sql"
-	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -27,33 +26,6 @@ import (
 	"github.com/coder/coder/v2/testutil"
 )
 
-type statsBatcher struct {
-	mu sync.Mutex
-
-	called          int64
-	lastTime        time.Time
-	lastAgentID     uuid.UUID
-	lastTemplateID  uuid.UUID
-	lastUserID      uuid.UUID
-	lastWorkspaceID uuid.UUID
-	lastStats       *agentproto.Stats
-}
-
-var _ agentapi.StatsBatcher = &statsBatcher{}
-
-func (b *statsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-	b.called++
-	b.lastTime = now
-	b.lastAgentID = agentID
-	b.lastTemplateID = templateID
-	b.lastUserID = userID
-	b.lastWorkspaceID = workspaceID
-	b.lastStats = st
-	return nil
-}
-
 func TestUpdateStates(t *testing.T) {
 	t.Parallel()
 
@@ -94,7 +66,6 @@ func TestUpdateStates(t *testing.T) {
 					panic("not implemented")
 				},
 			}
-			batcher                    = &statsBatcher{}
 			updateAgentMetricsFnCalled = false
 
 			req = &agentproto.UpdateStatsRequest{
@@ -134,7 +105,6 @@ func TestUpdateStates(t *testing.T) {
 			StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
 				Database:              dbM,
 				Pubsub:                ps,
-				StatsBatcher:          batcher,
 				TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
 				UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
 					updateAgentMetricsFnCalled = true
@@ -174,6 +144,9 @@ func TestUpdateStates(t *testing.T) {
 		// User gets fetched to hit the UpdateAgentMetricsFn.
 		dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
 
+		// Agent stats get inserted.
+		dbM.EXPECT().InsertWorkspaceAgentStats(gomock.Any(), gomock.Any())
+
 		// Ensure that pubsub notifications are sent.
 		notifyDescription := make(chan []byte)
 		ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) {
@@ -187,16 +160,6 @@ func TestUpdateStates(t *testing.T) {
 		require.Equal(t, &agentproto.UpdateStatsResponse{
 			ReportInterval: durationpb.New(10 * time.Second),
 		}, resp)
-
-		batcher.mu.Lock()
-		defer batcher.mu.Unlock()
-		require.Equal(t, int64(1), batcher.called)
-		require.Equal(t, now, batcher.lastTime)
-		require.Equal(t, agent.ID, batcher.lastAgentID)
-		require.Equal(t, template.ID, batcher.lastTemplateID)
-		require.Equal(t, user.ID, batcher.lastUserID)
-		require.Equal(t, workspace.ID, batcher.lastWorkspaceID)
-		require.Equal(t, req.Stats, batcher.lastStats)
 		ctx := testutil.Context(t, testutil.WaitShort)
 		select {
 		case <-ctx.Done():
@@ -222,7 +185,6 @@ func TestUpdateStates(t *testing.T) {
 					panic("not implemented")
 				},
 			}
-			batcher = &statsBatcher{}
 
 			req = &agentproto.UpdateStatsRequest{
 				Stats: &agentproto.Stats{
@@ -240,7 +202,6 @@ func TestUpdateStates(t *testing.T) {
 			StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
 				Database:              dbM,
 				Pubsub:                ps,
-				StatsBatcher:          batcher,
 				TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
 				// Ignored when nil.
 				UpdateAgentMetricsFn: nil,
@@ -263,6 +224,9 @@ func TestUpdateStates(t *testing.T) {
 			LastUsedAt: now,
 		}).Return(nil)
 
+		// Agent stats get inserted.
+		dbM.EXPECT().InsertWorkspaceAgentStats(gomock.Any(), gomock.Any())
+
 		_, err := api.UpdateStats(context.Background(), req)
 		require.NoError(t, err)
 	})
@@ -285,7 +249,6 @@ func TestUpdateStates(t *testing.T) {
 			StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
 				Database:              dbM,
 				Pubsub:                ps,
-				StatsBatcher:          nil, // should not be called
 				TemplateScheduleStore: nil, // should not be called
 				UpdateAgentMetricsFn:  nil, // should not be called
 			}),
@@ -336,7 +299,6 @@ func TestUpdateStates(t *testing.T) {
 					panic("not implemented")
 				},
 			}
-			batcher                    = &statsBatcher{}
 			updateAgentMetricsFnCalled = false
 
 			req = &agentproto.UpdateStatsRequest{
@@ -357,7 +319,6 @@ func TestUpdateStates(t *testing.T) {
 			StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
 				Database:              dbM,
 				Pubsub:                ps,
-				StatsBatcher:          batcher,
 				TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
 				UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
 					updateAgentMetricsFnCalled = true
@@ -398,6 +359,9 @@ func TestUpdateStates(t *testing.T) {
 		// User gets fetched to hit the UpdateAgentMetricsFn.
 		dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
 
+		// Agent stats get inserted.
+		dbM.EXPECT().InsertWorkspaceAgentStats(gomock.Any(), gomock.Any())
+
 		resp, err := api.UpdateStats(context.Background(), req)
 		require.NoError(t, err)
 		require.Equal(t, &agentproto.UpdateStatsResponse{

diff --git a/coderd/apikey/apikey_test.go b/coderd/apikey/apikey_test.go
@@ -128,7 +128,7 @@ func TestGenerate(t *testing.T) {
 
 			// Assert that the hashed secret is correct.
 			hashed := sha256.Sum256([]byte(keytokens[1]))
-			assert.ElementsMatch(t, hashed, key.HashedSecret[:])
+			assert.ElementsMatch(t, hashed, key.HashedSecret)
 
 			assert.Equal(t, tc.params.UserID, key.UserID)
 			assert.WithinDuration(t, dbtime.Now(), key.CreatedAt, time.Second*5)