Push startup script logs from agent

code-asher · code-asher · commit 1cc3e9d674aa · 2023-03-01T08:19:41.000-09:00
diff --git a/agent/agent.go b/agent/agent.go
@@ -23,6 +23,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/ammario/prefixsuffix"
 	"github.com/armon/circbuf"
 	"github.com/gliderlabs/ssh"
 	"github.com/google/uuid"
@@ -76,6 +77,7 @@ type Client interface {
 	PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error
 	PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error
 	PostStartup(ctx context.Context, req agentsdk.PostStartupRequest) error
+	InsertOrUpdateStartupLogs(ctx context.Context, req agentsdk.InsertOrUpdateStartupLogsRequest) error
 }
 
 func New(options Options) io.Closer {
@@ -617,13 +619,26 @@ func (a *agent) runStartupScript(ctx context.Context, script string) error {
 	}
 
 	a.logger.Info(ctx, "running startup script", slog.F("script", script))
-	writer, err := a.filesystem.OpenFile(filepath.Join(a.logDir, "coder-startup-script.log"), os.O_CREATE|os.O_RDWR, 0o600)
+
+	fileWriter, err := a.filesystem.OpenFile(filepath.Join(a.logDir, "coder-startup-script.log"), os.O_CREATE|os.O_RDWR, 0o600)
 	if err != nil {
 		return xerrors.Errorf("open startup script log file: %w", err)
 	}
 	defer func() {
-		_ = writer.Close()
+		_ = fileWriter.Close()
+	}()
+
+	saver := &prefixsuffix.Saver{N: 512 << 10}
+	writer := io.MultiWriter(saver, fileWriter)
+	defer func() {
+		err := a.client.InsertOrUpdateStartupLogs(ctx, agentsdk.InsertOrUpdateStartupLogsRequest{
+			Output: string(saver.Bytes()),
+		})
+		if err != nil {
+			a.logger.Error(ctx, "upload startup logs", slog.Error(err))
+		}
 	}()
+
 	cmd, err := a.createCommand(ctx, script, nil)
 	if err != nil {
 		return xerrors.Errorf("create command: %w", err)
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -667,9 +667,9 @@ func TestAgent_StartupScript(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("This test doesn't work on Windows for some reason...")
 	}
-	content := "output"
+	content := "output\n"
 	//nolint:dogsled
-	_, _, _, fs := setupAgent(t, agentsdk.Metadata{
+	_, client, _, fs := setupAgent(t, agentsdk.Metadata{
 		StartupScript: "echo " + content,
 	}, 0)
 	var gotContent string
@@ -694,7 +694,8 @@ func TestAgent_StartupScript(t *testing.T) {
 		gotContent = string(content)
 		return true
 	}, testutil.WaitShort, testutil.IntervalMedium)
-	require.Equal(t, content, strings.TrimSpace(gotContent))
+	require.Equal(t, content, gotContent)
+	require.Equal(t, content, client.getLogs())
 }
 
 func TestAgent_Lifecycle(t *testing.T) {
@@ -1229,6 +1230,7 @@ type client struct {
 	mu              sync.Mutex // Protects following.
 	lifecycleStates []codersdk.WorkspaceAgentLifecycle
 	startup         agentsdk.PostStartupRequest
+	logs            agentsdk.InsertOrUpdateStartupLogsRequest
 }
 
 func (c *client) Metadata(_ context.Context) (agentsdk.Metadata, error) {
@@ -1313,6 +1315,19 @@ func (c *client) PostStartup(_ context.Context, startup agentsdk.PostStartupRequ
 	return nil
 }
 
+func (c *client) getLogs() string {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.logs.Output
+}
+
+func (c *client) InsertOrUpdateStartupLogs(_ context.Context, logs agentsdk.InsertOrUpdateStartupLogsRequest) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.logs = logs
+	return nil
+}
+
 // tempDirUnixSocket returns a temporary directory that can safely hold unix
 // sockets (probably).
 //
diff --git a/go.mod b/go.mod
@@ -59,6 +59,7 @@ require (
 	cloud.google.com/go/compute/metadata v0.2.1
 	github.com/AlecAivazis/survey/v2 v2.3.5
 	github.com/adrg/xdg v0.4.0
+	github.com/ammario/prefixsuffix v0.0.0-20200405191514-5a0456bf2cfd
 	github.com/andybalholm/brotli v1.0.4
 	github.com/armon/circbuf v0.0.0-20190214190532-5111143e8da2
 	github.com/awalterschulze/gographviz v2.0.3+incompatible
diff --git a/go.sum b/go.sum
@@ -190,6 +190,8 @@ github.com/alexbrainman/sspi v0.0.0-20210105120005-909beea2cc74/go.mod h1:cEWa1L
 github.com/alexflint/go-filemutex v0.0.0-20171022225611-72bdc8eae2ae/go.mod h1:CgnQgUtFrFz9mxFNtED3jI5tLDjKlOM+oUF/sTk6ps0=
 github.com/alexflint/go-filemutex v1.1.0/go.mod h1:7P4iRhttt/nUvUOrYIhcpMzv2G6CY9UnI16Z+UJqRyk=
 github.com/alexkohler/prealloc v1.0.0/go.mod h1:VetnK3dIgFBBKmg0YnD9F9x6Icjd+9cvfHR56wJVlKE=
+github.com/ammario/prefixsuffix v0.0.0-20200405191514-5a0456bf2cfd h1:WOzjyD34+0vVw3wzE7js8Yvzo08ljzvK1jG6wL8elVU=
+github.com/ammario/prefixsuffix v0.0.0-20200405191514-5a0456bf2cfd/go.mod h1:VM1c/0Tl3O26UkHMbU32VFqLwLvi2FA40b6s5vPOpoo=
 github.com/andreyvit/diff v0.0.0-20170406064948-c7f18ee00883 h1:bvNMNQO63//z+xNgfBlViaCIJKLlCJ6/fmUseuG0wVQ=
 github.com/andybalholm/brotli v1.0.2/go.mod h1:loMXtMfwqflxFJPmdbJO0a3KNoPuLBgiu3qAvBg8x/Y=
 github.com/andybalholm/brotli v1.0.3/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
