Skip to content

perf: remove pubsub usage for workspace agent metadata updates #21177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
cstyan wants to merge 1 commit into from
Closed

perf: remove pubsub usage for workspace agent metadata updates #21177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion coderd/agentapi/api.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,6 @@ func New(opts Options, workspace database.Workspace) *API {
AgentFn: api.agent,
Workspace: api.cachedWorkspaceFields,
Database: opts.Database,
Pubsub: opts.Pubsub,
Log: opts.Log,
}

Expand Down
15 changes: 0 additions & 15 deletions coderd/agentapi/metadata.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package agentapi

import (
"context"
"encoding/json"
"fmt"
"time"

Expand All @@ -14,14 +13,12 @@ import (
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/database/pubsub"
)

type MetadataAPI struct {
AgentFn func(context.Context) (database.WorkspaceAgent, error)
Workspace *CachedWorkspaceFields
Database database.Store
Pubsub pubsub.Pubsub
Log slog.Logger

TimeNowFn func() time.Time // defaults to dbtime.Now()
Expand Down Expand Up @@ -126,18 +123,6 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
return nil, xerrors.Errorf("update workspace agent metadata in database: %w", err)
}

payload, err := json.Marshal(WorkspaceAgentMetadataChannelPayload{
CollectedAt: collectedAt,
Keys: dbUpdate.Key,
})
if err != nil {
return nil, xerrors.Errorf("marshal workspace agent metadata channel payload: %w", err)
}
err = a.Pubsub.Publish(WatchWorkspaceAgentMetadataChannel(workspaceAgent.ID), payload)
if err != nil {
return nil, xerrors.Errorf("publish workspace agent metadata: %w", err)
}

// If the metadata keys were too large, we return an error so the agent can
// log it.
if allKeysLen > maxAllKeysLen {
Expand Down
58 changes: 0 additions & 58 deletions coderd/agentapi/metadata_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ package agentapi_test
import (
"context"
"database/sql"
"encoding/json"
"sync/atomic"
"testing"
"time"
Expand All @@ -20,26 +19,12 @@ import (
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbmock"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/database/pubsub"
"github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/coderd/rbac/policy"
"github.com/coder/coder/v2/testutil"
"github.com/coder/quartz"
)

type fakePublisher struct {
// Nil pointer to pass interface check.
pubsub.Pubsub
publishes [][]byte
}

var _ pubsub.Pubsub = &fakePublisher{}

func (f *fakePublisher) Publish(_ string, message []byte) error {
f.publishes = append(f.publishes, message)
return nil
}

func TestBatchUpdateMetadata(t *testing.T) {
t.Parallel()

Expand All @@ -51,7 +36,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
t.Parallel()

dbM := dbmock.NewMockStore(gomock.NewController(t))
pub := &fakePublisher{}

now := dbtime.Now()
req := &agentproto.BatchUpdateMetadataRequest{
Expand Down Expand Up @@ -92,7 +76,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
},
Workspace: &agentapi.CachedWorkspaceFields{},
Database: dbM,
Pubsub: pub,
Log: testutil.Logger(t),
TimeNowFn: func() time.Time {
return now
Expand All @@ -102,21 +85,12 @@ func TestBatchUpdateMetadata(t *testing.T) {
resp, err := api.BatchUpdateMetadata(context.Background(), req)
require.NoError(t, err)
require.Equal(t, &agentproto.BatchUpdateMetadataResponse{}, resp)

require.Equal(t, 1, len(pub.publishes))
var gotEvent agentapi.WorkspaceAgentMetadataChannelPayload
require.NoError(t, json.Unmarshal(pub.publishes[0], &gotEvent))
require.Equal(t, agentapi.WorkspaceAgentMetadataChannelPayload{
CollectedAt: now,
Keys: []string{req.Metadata[0].Key, req.Metadata[1].Key},
}, gotEvent)
})

t.Run("ExceededLength", func(t *testing.T) {
t.Parallel()

dbM := dbmock.NewMockStore(gomock.NewController(t))
pub := pubsub.NewInMemory()

almostLongValue := ""
for i := 0; i < 2048; i++ {
Expand Down Expand Up @@ -178,7 +152,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
},
Workspace: &agentapi.CachedWorkspaceFields{},
Database: dbM,
Pubsub: pub,
Log: testutil.Logger(t),
TimeNowFn: func() time.Time {
return now
Expand All @@ -194,7 +167,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
t.Parallel()

dbM := dbmock.NewMockStore(gomock.NewController(t))
pub := pubsub.NewInMemory()

now := dbtime.Now()
req := &agentproto.BatchUpdateMetadataRequest{
Expand Down Expand Up @@ -248,38 +220,16 @@ func TestBatchUpdateMetadata(t *testing.T) {
},
Workspace: &agentapi.CachedWorkspaceFields{},
Database: dbM,
Pubsub: pub,
Log: testutil.Logger(t),
TimeNowFn: func() time.Time {
return now
},
}

// Watch the pubsub for events.
var (
eventCount int64
gotEvent agentapi.WorkspaceAgentMetadataChannelPayload
)
cancel, err := pub.Subscribe(agentapi.WatchWorkspaceAgentMetadataChannel(agent.ID), func(ctx context.Context, message []byte) {
if atomic.AddInt64(&eventCount, 1) > 1 {
return
}
require.NoError(t, json.Unmarshal(message, &gotEvent))
})
require.NoError(t, err)
defer cancel()

resp, err := api.BatchUpdateMetadata(context.Background(), req)
require.Error(t, err)
require.Equal(t, "metadata keys of 6145 bytes exceeded 6144 bytes", err.Error())
require.Nil(t, resp)

require.Equal(t, int64(1), atomic.LoadInt64(&eventCount))
require.Equal(t, agentapi.WorkspaceAgentMetadataChannelPayload{
CollectedAt: now,
// No key 4.
Keys: []string{req.Metadata[0].Key, req.Metadata[1].Key, req.Metadata[2].Key},
}, gotEvent)
})

// Test RBAC fast path with valid RBAC object - should NOT call GetWorkspaceByAgentID
Expand All @@ -291,7 +241,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
var (
ctrl = gomock.NewController(t)
dbM = dbmock.NewMockStore(ctrl)
pub = &fakePublisher{}
now = dbtime.Now()
// Set up consistent IDs that represent a valid workspace->agent relationship
workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
Expand Down Expand Up @@ -345,7 +294,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
},
Workspace: &agentapi.CachedWorkspaceFields{},
Database: dbauthz.New(dbM, auth, testutil.Logger(t), accessControlStore),
Pubsub: pub,
Log: testutil.Logger(t),
TimeNowFn: func() time.Time {
return now
Expand Down Expand Up @@ -373,7 +321,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
var (
ctrl = gomock.NewController(t)
dbM = dbmock.NewMockStore(ctrl)
pub = &fakePublisher{}
now = dbtime.Now()
workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321")
Expand Down Expand Up @@ -430,7 +377,6 @@ func TestBatchUpdateMetadata(t *testing.T) {

Workspace: &agentapi.CachedWorkspaceFields{},
Database: dbauthz.New(dbM, auth, testutil.Logger(t), accessControlStore),
Pubsub: pub,
Log: testutil.Logger(t),
TimeNowFn: func() time.Time {
return now
Expand Down Expand Up @@ -460,7 +406,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
var (
ctrl = gomock.NewController(t)
dbM = dbmock.NewMockStore(ctrl)
pub = &fakePublisher{}
now = dbtime.Now()
workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321")
Expand Down Expand Up @@ -516,7 +461,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
},
Workspace: &agentapi.CachedWorkspaceFields{},
Database: dbauthz.New(dbM, auth, testutil.Logger(t), accessControlStore),
Pubsub: pub,
Log: testutil.Logger(t),
TimeNowFn: func() time.Time {
return now
Expand All @@ -539,7 +483,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
var (
ctrl = gomock.NewController(t)
dbM = dbmock.NewMockStore(ctrl)
pub = &fakePublisher{}
now = dbtime.Now()
mClock = quartz.NewMock(t)
tickerTrap = mClock.Trap().TickerFunc("cache_refresh")
Expand Down Expand Up @@ -679,7 +622,6 @@ func TestBatchUpdateMetadata(t *testing.T) {
Database: dbauthz.New(dbM, auth, testutil.Logger(t), accessControlStore),
Log: testutil.Logger(t),
Clock: mClock,
Pubsub: pub,
}, initialWorkspace) // Cache is initialized with 9am schedule and "my-workspace" name

// Wait for ticker to be set up and release it so it can fire
Expand Down
1 change: 1 addition & 0 deletions coderd/coderd.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1428,6 +1428,7 @@ func New(options *Options) *API {
httpmw.ExtractWorkspaceParam(options.Database),
)
r.Get("/", api.workspaceAgent)
r.Get("/metadata", api.workspaceAgentMetadata)
r.Get("/watch-metadata", api.watchWorkspaceAgentMetadataSSE)
r.Get("/watch-metadata-ws", api.watchWorkspaceAgentMetadataWS)
r.Get("/startup-logs", api.workspaceAgentLogsDeprecated)
Expand Down
26 changes: 26 additions & 0 deletions coderd/workspaceagents.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1594,6 +1594,32 @@ func convertScripts(dbScripts []database.WorkspaceAgentScript) []codersdk.Worksp
return scripts
}

// @Summary Get workspace agent metadata
// @ID get-workspace-agent-metadata
// @Security CoderSessionToken
// @Produce json
// @Tags Agents
// @Success 200 {array} codersdk.WorkspaceAgentMetadata
// @Param workspaceagent path string true "Workspace agent ID" format(uuid)
// @Router /workspaceagents/{workspaceagent}/metadata [get]
func (api *API) workspaceAgentMetadata(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
workspaceAgent := httpmw.WorkspaceAgentParam(r)

metadata, err := api.Database.GetWorkspaceAgentMetadata(ctx,
database.GetWorkspaceAgentMetadataParams{
WorkspaceAgentID: workspaceAgent.ID,
Keys: nil,
})
if err != nil {
httpapi.InternalServerError(rw, err)
return
}

httpapi.Write(ctx, rw, http.StatusOK,
convertWorkspaceAgentMetadata(metadata))
}

// @Summary Watch for workspace agent metadata updates
// @ID watch-for-workspace-agent-metadata-updates
// @Security CoderSessionToken
Expand Down
10 changes: 10 additions & 0 deletions site/src/api/api.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1647,6 +1647,16 @@ class ApiMethods {
return response.data;
};

getWorkspaceAgentMetadata = async (
agentID: string,
): Promise<TypesGen.WorkspaceAgentMetadata[]> => {
const response = await this.axios.get<TypesGen.WorkspaceAgentMetadata[]>(
`/api/v2/workspaceagents/${agentID}/metadata`,
);

return response.data;
};

putWorkspaceExtension = async (
workspaceId: string,
newDeadline: dayjs.Dayjs,
Expand Down
Loading
Loading