Skip to content

Commit 24f6614

Browse files
deansheatherdeansheather
committed
Remaining tests
1 parent 36bfcbd commit 24f6614

File tree

9 files changed

+954
-71
lines changed

9 files changed

+954
-71
lines changed

‎coderd/agentapi/activitybump.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,13 +41,14 @@ func ActivityBumpWorkspace(ctx context.Context, log slog.Logger, db database.Sto
4141
// low priority operations fail first.
4242
ctx, cancel := context.WithTimeout(ctx, time.Second*15)
4343
defer cancel()
44-
if err := db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{
44+
err := db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{
4545
NextAutostart: nextAutostart.UTC(),
4646
WorkspaceID: workspaceID,
47-
}); err != nil {
47+
})
48+
if err != nil {
4849
if !xerrors.Is(err, context.Canceled) && !database.IsQueryCanceledError(err) {
4950
// Bump will fail if the context is canceled, but this is ok.
50-
log.Error(ctx, "bump failed", slog.Error(err),
51+
log.Error(ctx, "activity bump failed", slog.Error(err),
5152
slog.F("workspace_id", workspaceID),
5253
)
5354
}

‎coderd/agentapi/api.go

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ import (
1717

1818
"cdr.dev/slog"
1919
agentproto "github.com/coder/coder/v2/agent/proto"
20-
"github.com/coder/coder/v2/coderd/batchstats"
2120
"github.com/coder/coder/v2/coderd/database"
2221
"github.com/coder/coder/v2/coderd/database/pubsub"
2322
"github.com/coder/coder/v2/coderd/externalauth"
@@ -58,7 +57,7 @@ type Options struct {
5857
Pubsub pubsub.Pubsub
5958
DerpMapFn func() *tailcfg.DERPMap
6059
TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
61-
StatsBatcher *batchstats.Batcher
60+
StatsBatcher StatsBatcher // *batchstats.Batcher
6261
PublishWorkspaceUpdateFn func(ctx context.Context, workspaceID uuid.UUID)
6362
PublishWorkspaceAgentLogsUpdateFn func(ctx context.Context, workspaceAgentID uuid.UUID, msg agentsdk.LogsNotifyMessage)
6463

@@ -201,20 +200,15 @@ func (a *API) workspaceID(ctx context.Context, agent *database.WorkspaceAgent) (
201200
agent = &agnt
202201
}
203202

204-
resource, err := a.opts.Database.GetWorkspaceResourceByID(ctx, agent.ResourceID)
203+
getWorkspaceAgentByIDRow, err := a.opts.Database.GetWorkspaceByAgentID(ctx, agent.ID)
205204
if err != nil {
206-
return uuid.Nil, xerrors.Errorf("get workspace agent resource by id %q: %w", agent.ResourceID, err)
207-
}
208-
209-
build, err := a.opts.Database.GetWorkspaceBuildByJobID(ctx, resource.JobID)
210-
if err != nil {
211-
return uuid.Nil, xerrors.Errorf("get workspace build by job id %q: %w", resource.JobID, err)
205+
return uuid.Nil, xerrors.Errorf("get workspace by agent id %q: %w", agent.ID, err)
212206
}
213207

214208
a.mu.Lock()
215-
a.cachedWorkspaceID = build.WorkspaceID
209+
a.cachedWorkspaceID = getWorkspaceAgentByIDRow.Workspace.ID
216210
a.mu.Unlock()
217-
return build.WorkspaceID, nil
211+
return getWorkspaceAgentByIDRow.Workspace.ID, nil
218212
}
219213

220214
func (a *API) publishWorkspaceUpdate(ctx context.Context, agent *database.WorkspaceAgent) error {

‎coderd/agentapi/metadata.go

Lines changed: 55 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"cdr.dev/slog"
1313
agentproto "github.com/coder/coder/v2/agent/proto"
1414
"github.com/coder/coder/v2/coderd/database"
15+
"github.com/coder/coder/v2/coderd/database/dbtime"
1516
"github.com/coder/coder/v2/coderd/database/pubsub"
1617
)
1718

@@ -20,14 +21,26 @@ type MetadataAPI struct {
2021
Database database.Store
2122
Pubsub pubsub.Pubsub
2223
Log slog.Logger
24+
25+
TimeNowFn func() time.Time // defaults to dbtime.Now()
26+
}
27+
28+
func (a *MetadataAPI) now() time.Time {
29+
if a.TimeNowFn != nil {
30+
return a.TimeNowFn()
31+
}
32+
return dbtime.Now()
2333
}
2434

2535
func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.BatchUpdateMetadataRequest) (*agentproto.BatchUpdateMetadataResponse, error) {
2636
const (
27-
// maxValueLen is set to 2048 to stay under the 8000 byte Postgres
28-
// NOTIFY limit. Since both value and error can be set, the real payload
29-
// limit is 2 * 2048 * 4/3 <base64 expansion> = 5461 bytes + a few
30-
// hundred bytes for JSON syntax, key names, and metadata.
37+
// maxAllKeysLen is the maximum length of all metadata keys. This is
38+
// 6144 to stay below the Postgres NOTIFY limit of 8000 bytes, with some
39+
// headway for the timestamp and JSON encoding. Any values that would
40+
// exceed this limit are discarded (the rest are still inserted) and an
41+
// error is returned.
42+
maxAllKeysLen = 6144 // 1024 * 6
43+
3144
maxValueLen = 2048
3245
maxErrorLen = maxValueLen
3346
)
@@ -37,18 +50,36 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
3750
return nil, err
3851
}
3952

40-
collectedAt := time.Now()
41-
dbUpdate := database.UpdateWorkspaceAgentMetadataParams{
42-
WorkspaceAgentID: workspaceAgent.ID,
43-
Key: make([]string, 0, len(req.Metadata)),
44-
Value: make([]string, 0, len(req.Metadata)),
45-
Error: make([]string, 0, len(req.Metadata)),
46-
CollectedAt: make([]time.Time, 0, len(req.Metadata)),
47-
}
48-
53+
var (
54+
collectedAt = a.now()
55+
allKeysLen = 0
56+
dbUpdate = database.UpdateWorkspaceAgentMetadataParams{
57+
WorkspaceAgentID: workspaceAgent.ID,
58+
// These need to be `make(x, 0, len(req.Metadata))` instead of
59+
// `make(x, len(req.Metadata))` because we may not insert all
60+
// metadata if the keys are large.
61+
Key: make([]string, 0, len(req.Metadata)),
62+
Value: make([]string, 0, len(req.Metadata)),
63+
Error: make([]string, 0, len(req.Metadata)),
64+
CollectedAt: make([]time.Time, 0, len(req.Metadata)),
65+
}
66+
)
4967
for _, md := range req.Metadata {
5068
metadataError := md.Result.Error
5169

70+
allKeysLen += len(md.Key)
71+
if allKeysLen > maxAllKeysLen {
72+
// We still insert the rest of the metadata, and we return an error
73+
// after the insert.
74+
a.Log.Warn(
75+
ctx, "discarded extra agent metadata due to excessive key length",
76+
slog.F("collected_at", collectedAt),
77+
slog.F("all_keys_len", allKeysLen),
78+
slog.F("max_all_keys_len", maxAllKeysLen),
79+
)
80+
break
81+
}
82+
5283
// We overwrite the error if the provided payload is too long.
5384
if len(md.Result.Value) > maxValueLen {
5485
metadataError = fmt.Sprintf("value of %d bytes exceeded %d bytes", len(md.Result.Value), maxValueLen)
@@ -71,30 +102,34 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
71102
a.Log.Debug(
72103
ctx, "accepted metadata report",
73104
slog.F("collected_at", collectedAt),
74-
slog.F("original_collected_at", collectedAt),
75105
slog.F("key", md.Key),
76106
slog.F("value", ellipse(md.Result.Value, 16)),
77107
)
78108
}
79109

110+
err = a.Database.UpdateWorkspaceAgentMetadata(ctx, dbUpdate)
111+
if err != nil {
112+
return nil, xerrors.Errorf("update workspace agent metadata in database: %w", err)
113+
}
114+
80115
payload, err := json.Marshal(WorkspaceAgentMetadataChannelPayload{
81116
CollectedAt: collectedAt,
82117
Keys: dbUpdate.Key,
83118
})
84119
if err != nil {
85120
return nil, xerrors.Errorf("marshal workspace agent metadata channel payload: %w", err)
86121
}
87-
88-
err = a.Database.UpdateWorkspaceAgentMetadata(ctx, dbUpdate)
89-
if err != nil {
90-
return nil, xerrors.Errorf("update workspace agent metadata in database: %w", err)
91-
}
92-
93122
err = a.Pubsub.Publish(WatchWorkspaceAgentMetadataChannel(workspaceAgent.ID), payload)
94123
if err != nil {
95124
return nil, xerrors.Errorf("publish workspace agent metadata: %w", err)
96125
}
97126

127+
// If the metadata keys were too large, we return an error so the agent can
128+
// log it.
129+
if allKeysLen > maxAllKeysLen {
130+
return nil, xerrors.Errorf("metadata keys of %d bytes exceeded %d bytes", allKeysLen, maxAllKeysLen)
131+
}
132+
98133
return &agentproto.BatchUpdateMetadataResponse{}, nil
99134
}
100135

0 commit comments

Comments
 (0)