Add tests for streaming logs

coder · kylecarbs · Mar 23, 2023 · Feb 23, 2023 · Feb 23, 2023 · Feb 23, 2023
commit 736705f87068ded10740972bfdaeefb31e46f24c
diff --git a/agent/agent.go b/agent/agent.go
@@ -88,7 +88,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
-	AppendStartupLogs(ctx context.Context, req []agentsdk.StartupLog) error
+	PatchStartupLogs(ctx context.Context, req agentsdk.PatchStartupLogs) error
 }
 
 func New(options Options) io.Closer {
@@ -202,19 +202,6 @@ func (a *agent) runLoop(ctx context.Context) {
 	}
 }
 
-func (a *agent) appendStartupLogsLoop(ctx context.Context, logs []agentsdk.StartupLog) {
-	for r := retry.New(time.Second, 15*time.Second); r.Wait(ctx); {
-		err := a.client.AppendStartupLogs(ctx, logs)
-		if err == nil {
-			return
-		}
-		if errors.Is(err, context.Canceled) || xerrors.Is(err, context.DeadlineExceeded) {
-			return
-		}
-		a.logger.Error(ctx, "failed to append startup logs", slog.Error(err))
-	}
-}
-
 // reportLifecycleLoop reports the current lifecycle state once.
 // Only the latest state is reported, intermediate states may be
 // lost if the agent can't communicate with the API.
@@ -680,25 +667,20 @@ func (a *agent) runScript(ctx context.Context, lifecycle, script string) error {
 			return
 		}
 		logsSending = true
-		toSend := make([]agentsdk.StartupLog, len(queuedLogs))
-		copy(toSend, queuedLogs)
+		logsToSend := queuedLogs
+		queuedLogs = make([]agentsdk.StartupLog, 0)
 		logMutex.Unlock()
 		for r := retry.New(time.Second, 5*time.Second); r.Wait(ctx); {
-			err := a.client.AppendStartupLogs(ctx, toSend)
+			err := a.client.PatchStartupLogs(ctx, agentsdk.PatchStartupLogs{
+				Logs: logsToSend,
+			})
 			if err == nil {
 				break
 			}
-			a.logger.Error(ctx, "upload startup logs", slog.Error(err))
-		}
-		if ctx.Err() != nil {
-			logMutex.Lock()
-			logsSending = false
-			logMutex.Unlock()
-			return
+			a.logger.Error(ctx, "upload startup logs", slog.Error(err), slog.F("to_send", logsToSend))
 		}
 		logMutex.Lock()
 		logsSending = false
-		queuedLogs = queuedLogs[len(toSend):]
 		logMutex.Unlock()
 	}
 	queueLog := func(log agentsdk.StartupLog) {

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -1576,10 +1576,10 @@ func (c *client) getStartupLogs() []agentsdk.StartupLog {
 	return c.logs
 }
 
-func (c *client) AppendStartupLogs(_ context.Context, logs []agentsdk.StartupLog) error {
+func (c *client) PatchStartupLogs(_ context.Context, logs agentsdk.PatchStartupLogs) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	c.logs = append(c.logs, logs...)
+	c.logs = append(c.logs, logs.Logs...)
 	return nil
 }
 

diff --git a/coderd/provisionerjobs_test.go b/coderd/provisionerjobs_test.go
@@ -89,36 +89,4 @@ func TestProvisionerJobLogs(t *testing.T) {
 			}
 		}
 	})
-
-	t.Run("List", func(t *testing.T) {
-		t.Parallel()
-		client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
-		user := coderdtest.CreateFirstUser(t, client)
-		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
-			Parse: echo.ParseComplete,
-			ProvisionApply: []*proto.Provision_Response{{
-				Type: &proto.Provision_Response_Log{
-					Log: &proto.Log{
-						Level:  proto.LogLevel_INFO,
-						Output: "log-output",
-					},
-				},
-			}, {
-				Type: &proto.Provision_Response_Complete{
-					Complete: &proto.Provision_Complete{},
-				},
-			}},
-		})
-		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
-		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
-		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
-		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
-
-		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
-		defer cancel()
-
-		logs, err := client.WorkspaceBuildLogsBefore(ctx, workspace.LatestBuild.ID, 0)
-		require.NoError(t, err)
-		require.Greater(t, len(logs), 1)
-	})
 }
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -226,23 +226,23 @@ func (api *API) postWorkspaceAgentStartup(rw http.ResponseWriter, r *http.Reques
 // @Accept json
 // @Produce json
 // @Tags Agents
-// @Param request body []agentsdk.StartupLog true "Startup logs"
+// @Param request body agentsdk.PatchStartupLogs true "Startup logs"
 // @Success 200
 // @Router /workspaceagents/me/startup-logs [patch]
 // @x-apidocgen {"skip": true}
 func (api *API) patchWorkspaceAgentStartupLogs(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	workspaceAgent := httpmw.WorkspaceAgent(r)
 
-	var req []agentsdk.StartupLog
+	var req agentsdk.PatchStartupLogs
 	if !httpapi.Read(ctx, rw, r, &req) {
 		return
 	}
 
 	createdAt := make([]time.Time, 0)
 	output := make([]string, 0)
 	outputLength := 0
-	for _, log := range req {
+	for _, log := range req.Logs {
 		createdAt = append(createdAt, log.CreatedAt)
 		output = append(output, log.Output)
 		outputLength += len(log.Output)
@@ -311,7 +311,7 @@ func (api *API) workspaceAgentStartupLogs(rw http.ResponseWriter, r *http.Reques
 	// This mostly copies how provisioner job logs are streamed!
 	var (
 		ctx      = r.Context()
-		agent    = httpmw.WorkspaceAgent(r)
+		agent    = httpmw.WorkspaceAgentParam(r)
 		logger   = api.Logger.With(slog.F("workspace_agent_id", agent.ID))
 		follow   = r.URL.Query().Has("follow")
 		afterRaw = r.URL.Query().Get("after")

diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -175,6 +175,65 @@ func TestWorkspaceAgent(t *testing.T) {
 	})
 }
 
+func TestWorkspaceAgentStartup(t *testing.T) {
+	t.Parallel()
+	ctx, cancelFunc := testutil.Context(t)
+	defer cancelFunc()
+	client := coderdtest.New(t, &coderdtest.Options{
+		IncludeProvisionerDaemon: true,
+	})
+	user := coderdtest.CreateFirstUser(t, client)
+	authToken := uuid.NewString()
+	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+		Parse:         echo.ParseComplete,
+		ProvisionPlan: echo.ProvisionComplete,
+		ProvisionApply: []*proto.Provision_Response{{
+			Type: &proto.Provision_Response_Complete{
+				Complete: &proto.Provision_Complete{
+					Resources: []*proto.Resource{{
+						Name: "example",
+						Type: "aws_instance",
+						Agents: []*proto.Agent{{
+							Id: uuid.NewString(),
+							Auth: &proto.Agent_Token{
+								Token: authToken,
+							},
+						}},
+					}},
+				},
+			},
+		}},
+	})
+	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+	coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+	workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+	build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+
+	agentClient := agentsdk.New(client.URL)
+	agentClient.SetSessionToken(authToken)
+	err := agentClient.PatchStartupLogs(ctx, agentsdk.PatchStartupLogs{
+		Logs: []agentsdk.StartupLog{{
+			CreatedAt: database.Now(),
+			Output:    "testing",
+		}},
+	})
+	require.NoError(t, err)
+
+	logs, closer, err := client.WorkspaceAgentStartupLogsAfter(ctx, build.Resources[0].Agents[0].ID, 0)
+	require.NoError(t, err)
+	defer func() {
+		_ = closer.Close()
+	}()
+	var log codersdk.WorkspaceAgentStartupLog
+	select {
+	case <-ctx.Done():
+	case log = <-logs:
+	}
+	require.NoError(t, ctx.Err())
+	require.Equal(t, "testing", log.Output)
+	cancelFunc()
+}
+
 func TestWorkspaceAgentListen(t *testing.T) {
 	t.Parallel()
 

diff --git a/coderd/wsconncache/wsconncache_test.go b/coderd/wsconncache/wsconncache_test.go
@@ -250,6 +250,6 @@ func (*client) PostStartup(_ context.Context, _ agentsdk.PostStartupRequest) err
 	return nil
 }
 
-func (*client) AppendStartupLogs(_ context.Context, _ []agentsdk.StartupLog) error {
+func (*client) PatchStartupLogs(_ context.Context, _ agentsdk.PatchStartupLogs) error {
 	return nil
 }
diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -521,9 +521,13 @@ type StartupLog struct {
 	Output    string    `json:"output"`
 }
 
-// AppendStartupLogs writes log messages to the agent startup script.
+type PatchStartupLogs struct {
+	Logs []StartupLog `json:"logs"`
+}
+
+// PatchStartupLogs writes log messages to the agent startup script.
 // Log messages are limited to 1MB in total.
-func (c *Client) AppendStartupLogs(ctx context.Context, req []StartupLog) error {
+func (c *Client) PatchStartupLogs(ctx context.Context, req PatchStartupLogs) error {
 	res, err := c.SDK.Request(ctx, http.MethodPatch, "/api/v2/workspaceagents/me/startup-logs", req)
 	if err != nil {
 		return err

diff --git a/codersdk/provisionerdaemons.go b/codersdk/provisionerdaemons.go
@@ -9,8 +9,6 @@ import (
 	"net"
 	"net/http"
 	"net/http/cookiejar"
-	"net/url"
-	"strconv"
 	"time"
 
 	"github.com/google/uuid"
@@ -100,27 +98,6 @@ type ProvisionerJobLog struct {
 	Output    string    `json:"output"`
 }
 
-// provisionerJobLogsBefore provides log output that occurred before a time.
-// This is abstracted from a specific job type to provide consistency between
-// APIs. Logs is the only shared route between jobs.
-func (c *Client) provisionerJobLogsBefore(ctx context.Context, path string, before int64) ([]ProvisionerJobLog, error) {
-	values := url.Values{}
-	if before != 0 {
-		values["before"] = []string{strconv.FormatInt(before, 10)}
-	}
-	res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("%s?%s", path, values.Encode()), nil)
-	if err != nil {
-		return nil, err
-	}
-	if res.StatusCode != http.StatusOK {
-		defer res.Body.Close()
-		return nil, ReadBodyAsError(res)
-	}
-
-	var logs []ProvisionerJobLog
-	return logs, json.NewDecoder(res.Body).Decode(&logs)
-}
-
 // provisionerJobLogsAfter streams logs that occurred after a specific time.
 func (c *Client) provisionerJobLogsAfter(ctx context.Context, path string, after int64) (<-chan ProvisionerJobLog, io.Closer, error) {
 	afterQuery := ""

diff --git a/codersdk/templateversions.go b/codersdk/templateversions.go
@@ -186,11 +186,6 @@ func (c *Client) TemplateVersionVariables(ctx context.Context, version uuid.UUID
 	return variables, json.NewDecoder(res.Body).Decode(&variables)
 }
 
-// TemplateVersionLogsBefore returns logs that occurred before a specific log ID.
-func (c *Client) TemplateVersionLogsBefore(ctx context.Context, version uuid.UUID, before int64) ([]ProvisionerJobLog, error) {
-	return c.provisionerJobLogsBefore(ctx, fmt.Sprintf("/api/v2/templateversions/%s/logs", version), before)
-}
-
 // TemplateVersionLogsAfter streams logs for a template version that occurred after a specific log ID.
 func (c *Client) TemplateVersionLogsAfter(ctx context.Context, version uuid.UUID, after int64) (<-chan ProvisionerJobLog, io.Closer, error) {
 	return c.provisionerJobLogsAfter(ctx, fmt.Sprintf("/api/v2/templateversions/%s/logs", version), after)
@@ -253,12 +248,6 @@ func (c *Client) TemplateVersionDryRunResources(ctx context.Context, version, jo
 	return resources, json.NewDecoder(res.Body).Decode(&resources)
 }
 
-// TemplateVersionDryRunLogsBefore returns logs for a template version dry-run
-// that occurred before a specific log ID.
-func (c *Client) TemplateVersionDryRunLogsBefore(ctx context.Context, version, job uuid.UUID, before int64) ([]ProvisionerJobLog, error) {
-	return c.provisionerJobLogsBefore(ctx, fmt.Sprintf("/api/v2/templateversions/%s/dry-run/%s/logs", version, job), before)
-}
-
 // TemplateVersionDryRunLogsAfter streams logs for a template version dry-run
 // that occurred after a specific log ID.
 func (c *Client) TemplateVersionDryRunLogsAfter(ctx context.Context, version, job uuid.UUID, after int64) (<-chan ProvisionerJobLog, io.Closer, error) {

diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"net/http/cookiejar"
@@ -314,6 +315,65 @@ func (c *Client) WorkspaceAgentListeningPorts(ctx context.Context, agentID uuid.
 	return listeningPorts, json.NewDecoder(res.Body).Decode(&listeningPorts)
 }
 
+func (c *Client) WorkspaceAgentStartupLogsAfter(ctx context.Context, agentID uuid.UUID, after int64) (<-chan WorkspaceAgentStartupLog, io.Closer, error) {
+	afterQuery := ""
+	if after != 0 {
+		afterQuery = fmt.Sprintf("&after=%d", after)
+	}
+	followURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/startup-logs?follow%s", agentID, afterQuery))
+	if err != nil {
+		return nil, nil, err
+	}
+	jar, err := cookiejar.New(nil)
+	if err != nil {
+		return nil, nil, xerrors.Errorf("create cookie jar: %w", err)
+	}
+	jar.SetCookies(followURL, []*http.Cookie{{
+		Name:  SessionTokenCookie,
+		Value: c.SessionToken(),
+	}})
+	httpClient := &http.Client{
+		Jar:       jar,
+		Transport: c.HTTPClient.Transport,
+	}
+	conn, res, err := websocket.Dial(ctx, followURL.String(), &websocket.DialOptions{
+		HTTPClient:      httpClient,
+		CompressionMode: websocket.CompressionDisabled,
+	})
+	if err != nil {
+		if res == nil {
+			return nil, nil, err
+		}
+		return nil, nil, ReadBodyAsError(res)
+	}
+	logs := make(chan WorkspaceAgentStartupLog)
+	closed := make(chan struct{})
+	ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageText)
+	decoder := json.NewDecoder(wsNetConn)
+	go func() {
+		defer close(closed)
+		defer close(logs)
+		defer conn.Close(websocket.StatusGoingAway, "")
+		var log WorkspaceAgentStartupLog
+		for {
+			err = decoder.Decode(&log)
+			if err != nil {
+				return
+			}
+			select {
+			case <-ctx.Done():
+				return
+			case logs <- log:
+			}
+		}
+	}()
+	return logs, closeFunc(func() error {
+		_ = wsNetConn.Close()
+		<-closed
+		return nil
+	}), nil
+}
+
 // GitProvider is a constant that represents the
 // type of providers that are supported within Coder.
 type GitProvider string