Skip to content

Commit 05d536c

Browse files
kylecarbskylecarbs
committed
Add warning for overflow
1 parent 4c5b630 commit 05d536c

File tree

18 files changed

+412
-108
lines changed

18 files changed

+412
-108
lines changed

‎agent/agent.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -677,6 +677,13 @@ func (a *agent) runScript(ctx context.Context, lifecycle, script string) error {
677677
if err == nil {
678678
break
679679
}
680+
var sdkErr *codersdk.Error
681+
if errors.As(err, &sdkErr) {
682+
if sdkErr.StatusCode() == http.StatusRequestEntityTooLarge {
683+
a.logger.Warn(ctx, "startup logs too large, dropping logs")
684+
break
685+
}
686+
}
680687
a.logger.Error(ctx, "upload startup logs", slog.Error(err), slog.F("to_send", logsToSend))
681688
}
682689
logMutex.Lock()

‎agent/agent_test.go

Lines changed: 73 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import (
88
"fmt"
99
"io"
1010
"net"
11+
"net/http"
12+
"net/http/httptest"
1113
"net/netip"
1214
"os"
1315
"os/exec"
@@ -38,6 +40,7 @@ import (
3840
"cdr.dev/slog"
3941
"cdr.dev/slog/sloggers/slogtest"
4042
"github.com/coder/coder/agent"
43+
"github.com/coder/coder/coderd/httpapi"
4144
"github.com/coder/coder/codersdk"
4245
"github.com/coder/coder/codersdk/agentsdk"
4346
"github.com/coder/coder/pty/ptytest"
@@ -737,25 +740,78 @@ func TestAgent_SSHConnectionEnvVars(t *testing.T) {
737740

738741
func TestAgent_StartupScript(t *testing.T) {
739742
t.Parallel()
740-
if runtime.GOOS == "windows" {
741-
t.Skip("This test doesn't work on Windows for some reason...")
742-
}
743743
output := "something"
744744
command := "sh -c 'echo " + output + "'"
745745
if runtime.GOOS == "windows" {
746746
command = "cmd.exe /c echo " + output
747747
}
748-
//nolint:dogsled
749-
_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
750-
StartupScript: command,
751-
}, 0)
752-
assert.Eventually(t, func() bool {
753-
got := client.getLifecycleStates()
754-
return len(got) > 0 && got[len(got)-1] == codersdk.WorkspaceAgentLifecycleReady
755-
}, testutil.WaitShort, testutil.IntervalMedium)
748+
t.Run("Success", func(t *testing.T) {
749+
t.Parallel()
750+
client := &client{
751+
t: t,
752+
agentID: uuid.New(),
753+
metadata: agentsdk.Metadata{
754+
StartupScript: command,
755+
DERPMap: &tailcfg.DERPMap{},
756+
},
757+
statsChan: make(chan *agentsdk.Stats),
758+
coordinator: tailnet.NewCoordinator(),
759+
}
760+
closer := agent.New(agent.Options{
761+
Client: client,
762+
Filesystem: afero.NewMemMapFs(),
763+
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
764+
ReconnectingPTYTimeout: 0,
765+
})
766+
t.Cleanup(func() {
767+
_ = closer.Close()
768+
})
769+
assert.Eventually(t, func() bool {
770+
got := client.getLifecycleStates()
771+
return len(got) > 0 && got[len(got)-1] == codersdk.WorkspaceAgentLifecycleReady
772+
}, testutil.WaitShort, testutil.IntervalMedium)
756773

757-
require.Len(t, client.getStartupLogs(), 1)
758-
require.Equal(t, output, client.getStartupLogs()[0].Output)
774+
require.Len(t, client.getStartupLogs(), 1)
775+
require.Equal(t, output, client.getStartupLogs()[0].Output)
776+
})
777+
// This ensures that even when coderd sends back that the startup
778+
// script has written too many lines it will still succeed!
779+
t.Run("OverflowsAndSkips", func(t *testing.T) {
780+
t.Parallel()
781+
client := &client{
782+
t: t,
783+
agentID: uuid.New(),
784+
metadata: agentsdk.Metadata{
785+
StartupScript: command,
786+
DERPMap: &tailcfg.DERPMap{},
787+
},
788+
patchWorkspaceLogs: func() error {
789+
resp := httptest.NewRecorder()
790+
httpapi.Write(context.Background(), resp, http.StatusRequestEntityTooLarge, codersdk.Response{
791+
Message: "Too many lines!",
792+
})
793+
res := resp.Result()
794+
defer res.Body.Close()
795+
return codersdk.ReadBodyAsError(res)
796+
},
797+
statsChan: make(chan *agentsdk.Stats),
798+
coordinator: tailnet.NewCoordinator(),
799+
}
800+
closer := agent.New(agent.Options{
801+
Client: client,
802+
Filesystem: afero.NewMemMapFs(),
803+
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
804+
ReconnectingPTYTimeout: 0,
805+
})
806+
t.Cleanup(func() {
807+
_ = closer.Close()
808+
})
809+
assert.Eventually(t, func() bool {
810+
got := client.getLifecycleStates()
811+
return len(got) > 0 && got[len(got)-1] == codersdk.WorkspaceAgentLifecycleReady
812+
}, testutil.WaitShort, testutil.IntervalMedium)
813+
require.Len(t, client.getStartupLogs(), 0)
814+
})
759815
}
760816

761817
func TestAgent_Lifecycle(t *testing.T) {
@@ -1481,6 +1537,7 @@ type client struct {
14811537
statsChan chan *agentsdk.Stats
14821538
coordinator tailnet.Coordinator
14831539
lastWorkspaceAgent func()
1540+
patchWorkspaceLogs func() error
14841541

14851542
mu sync.Mutex // Protects following.
14861543
lifecycleStates []codersdk.WorkspaceAgentLifecycle
@@ -1579,6 +1636,9 @@ func (c *client) getStartupLogs() []agentsdk.StartupLog {
15791636
func (c *client) PatchStartupLogs(_ context.Context, logs agentsdk.PatchStartupLogs) error {
15801637
c.mu.Lock()
15811638
defer c.mu.Unlock()
1639+
if c.patchWorkspaceLogs != nil {
1640+
return c.patchWorkspaceLogs()
1641+
}
15821642
c.logs = append(c.logs, logs.Logs...)
15831643
return nil
15841644
}

‎coderd/apidoc/docs.go

Lines changed: 12 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/apidoc/swagger.json

Lines changed: 12 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbauthz/querier.go

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1237,6 +1237,24 @@ func (q *querier) UpdateWorkspaceAgentLifecycleStateByID(ctx context.Context, ar
12371237
return q.db.UpdateWorkspaceAgentLifecycleStateByID(ctx, arg)
12381238
}
12391239

1240+
func (q *querier) UpdateWorkspaceAgentStartupLogOverflowByID(ctx context.Context, arg database.UpdateWorkspaceAgentStartupLogOverflowByIDParams) error {
1241+
agent, err := q.db.GetWorkspaceAgentByID(ctx, arg.ID)
1242+
if err != nil {
1243+
return err
1244+
}
1245+
1246+
workspace, err := q.db.GetWorkspaceByAgentID(ctx, agent.ID)
1247+
if err != nil {
1248+
return err
1249+
}
1250+
1251+
if err := q.authorizeContext(ctx, rbac.ActionUpdate, workspace); err != nil {
1252+
return err
1253+
}
1254+
1255+
return q.db.UpdateWorkspaceAgentStartupLogOverflowByID(ctx, arg)
1256+
}
1257+
12401258
func (q *querier) UpdateWorkspaceAgentStartupByID(ctx context.Context, arg database.UpdateWorkspaceAgentStartupByIDParams) error {
12411259
agent, err := q.db.GetWorkspaceAgentByID(ctx, arg.ID)
12421260
if err != nil {

‎coderd/database/dbauthz/querier_test.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -995,6 +995,16 @@ func (s *MethodTestSuite) TestWorkspace() {
995995
LifecycleState: database.WorkspaceAgentLifecycleStateCreated,
996996
}).Asserts(ws, rbac.ActionUpdate).Returns()
997997
}))
998+
s.Run("UpdateWorkspaceAgentStartupLogOverflowByID", s.Subtest(func(db database.Store, check *expects) {
999+
ws := dbgen.Workspace(s.T(), db, database.Workspace{})
1000+
build := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{WorkspaceID: ws.ID, JobID: uuid.New()})
1001+
res := dbgen.WorkspaceResource(s.T(), db, database.WorkspaceResource{JobID: build.JobID})
1002+
agt := dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{ResourceID: res.ID})
1003+
check.Args(database.UpdateWorkspaceAgentStartupLogOverflowByIDParams{
1004+
ID: agt.ID,
1005+
StartupLogsOverflowed: true,
1006+
}).Asserts(ws, rbac.ActionUpdate).Returns()
1007+
}))
9981008
s.Run("UpdateWorkspaceAgentStartupByID", s.Subtest(func(db database.Store, check *expects) {
9991009
ws := dbgen.Workspace(s.T(), db, database.Workspace{})
10001010
build := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{WorkspaceID: ws.ID, JobID: uuid.New()})

‎coderd/database/dbfake/databasefake.go

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3511,7 +3511,7 @@ func (q *fakeQuerier) UpdateWorkspaceAgentStartupByID(_ context.Context, arg dat
35113511
return sql.ErrNoRows
35123512
}
35133513

3514-
func (q *fakeQuerier) GetWorkspaceAgentStartupLogsAfter(ctx context.Context, arg database.GetWorkspaceAgentStartupLogsAfterParams) ([]database.WorkspaceAgentStartupLog, error) {
3514+
func (q *fakeQuerier) GetWorkspaceAgentStartupLogsAfter(_ context.Context, arg database.GetWorkspaceAgentStartupLogsAfterParams) ([]database.WorkspaceAgentStartupLog, error) {
35153515
if err := validateDatabaseType(arg); err != nil {
35163516
return nil, err
35173517
}
@@ -3532,7 +3532,7 @@ func (q *fakeQuerier) GetWorkspaceAgentStartupLogsAfter(ctx context.Context, arg
35323532
return logs, nil
35333533
}
35343534

3535-
func (q *fakeQuerier) InsertWorkspaceAgentStartupLogs(ctx context.Context, arg database.InsertWorkspaceAgentStartupLogsParams) ([]database.WorkspaceAgentStartupLog, error) {
3535+
func (q *fakeQuerier) InsertWorkspaceAgentStartupLogs(_ context.Context, arg database.InsertWorkspaceAgentStartupLogsParams) ([]database.WorkspaceAgentStartupLog, error) {
35363536
if err := validateDatabaseType(arg); err != nil {
35373537
return nil, err
35383538
}
@@ -3545,6 +3545,7 @@ func (q *fakeQuerier) InsertWorkspaceAgentStartupLogs(ctx context.Context, arg d
35453545
if len(q.workspaceAgentLogs) > 0 {
35463546
id = q.workspaceAgentLogs[len(q.workspaceAgentLogs)-1].ID
35473547
}
3548+
outputLength := int32(0)
35483549
for index, output := range arg.Output {
35493550
id++
35503551
logs = append(logs, database.WorkspaceAgentStartupLog{
@@ -3553,6 +3554,22 @@ func (q *fakeQuerier) InsertWorkspaceAgentStartupLogs(ctx context.Context, arg d
35533554
CreatedAt: arg.CreatedAt[index],
35543555
Output: output,
35553556
})
3557+
outputLength += int32(len(output))
3558+
}
3559+
for index, agent := range q.workspaceAgents {
3560+
if agent.ID != arg.AgentID {
3561+
continue
3562+
}
3563+
// Greater than 1MB, same as the PostgreSQL constraint!
3564+
if agent.StartupLogsLength+outputLength > (1 << 20) {
3565+
return nil, &pq.Error{
3566+
Constraint: "max_startup_logs_length",
3567+
Table: "workspace_agents",
3568+
}
3569+
}
3570+
agent.StartupLogsLength += outputLength
3571+
q.workspaceAgents[index] = agent
3572+
break
35563573
}
35573574
q.workspaceAgentLogs = append(q.workspaceAgentLogs, logs...)
35583575
return logs, nil
@@ -4703,3 +4720,20 @@ func (q *fakeQuerier) UpdateWorkspaceAgentLifecycleStateByID(_ context.Context,
47034720
}
47044721
return sql.ErrNoRows
47054722
}
4723+
4724+
func (q *fakeQuerier) UpdateWorkspaceAgentStartupLogOverflowByID(_ context.Context, arg database.UpdateWorkspaceAgentStartupLogOverflowByIDParams) error {
4725+
if err := validateDatabaseType(arg); err != nil {
4726+
return err
4727+
}
4728+
4729+
q.mutex.Lock()
4730+
defer q.mutex.Unlock()
4731+
for i, agent := range q.workspaceAgents {
4732+
if agent.ID == arg.ID {
4733+
agent.StartupLogsOverflowed = arg.StartupLogsOverflowed
4734+
q.workspaceAgents[i] = agent
4735+
return nil
4736+
}
4737+
}
4738+
return sql.ErrNoRows
4739+
}

‎coderd/database/dump.sql

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000109_add_startup_logs.up.sql

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,7 @@ CREATE INDEX workspace_agent_startup_logs_id_agent_id_idx ON workspace_agent_sta
88

99
-- The maximum length of startup logs is 1MB per workspace agent.
1010
ALTER TABLE workspace_agents ADD COLUMN startup_logs_length integer NOT NULL DEFAULT 0 CONSTRAINT max_startup_logs_length CHECK (startup_logs_length <= 1048576);
11+
ALTER TABLE workspace_agents ADD COLUMN startup_logs_overflowed boolean NOT NULL DEFAULT false;
12+
1113
COMMENT ON COLUMN workspace_agents.startup_logs_length IS 'Total length of startup logs';
14+
COMMENT ON COLUMN workspace_agents.startup_logs_overflowed IS 'Whether the startup logs overflowed in length';

‎coderd/database/models.go

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier.go

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)