fix: avoid missed logs when streaming startup logs (#8029)

* feat(coderd,agent): send startup log eof at the end

* fix(coderd): fix edge case in startup log pubsub

* fix(coderd): ensure startup logs are closed on lifecycle state change (fallback)

* fix(codersdk): fix startup log channel shared memory bug

* fix(site): remove the EOF log line

coderd/database/dbauthz/dbauthz.go

@@ -1407,6 +1407,14 @@ func (q *querier) GetWorkspaceAgentStartupLogsAfter(ctx context.Context, arg dat
 	return q.db.GetWorkspaceAgentStartupLogsAfter(ctx, arg)
 }
 
+func (q *querier) GetWorkspaceAgentStartupLogsEOF(ctx context.Context, agentID uuid.UUID) (bool, error) {
+	_, err := q.GetWorkspaceAgentByID(ctx, agentID)
+	if err != nil {
+		return false, err
+	}
+	return q.db.GetWorkspaceAgentStartupLogsEOF(ctx, agentID)
+}
+
 func (q *querier) GetWorkspaceAgentStats(ctx context.Context, createdAfter time.Time) ([]database.GetWorkspaceAgentStatsRow, error) {
 	return q.db.GetWorkspaceAgentStats(ctx, createdAfter)
 }

coderd/database/dbfake/dbfake.go

@@ -2722,7 +2722,7 @@ func (q *fakeQuerier) GetWorkspaceAgentStartupLogsAfter(_ context.Context, arg d
 		if log.AgentID != arg.AgentID {
 			continue
 		}
-		if arg.CreatedAfter != 0 && log.ID < arg.CreatedAfter {
+		if arg.CreatedAfter != 0 && log.ID <= arg.CreatedAfter {
 			continue
 		}
 		logs = append(logs, log)
@@ -2730,6 +2730,22 @@ func (q *fakeQuerier) GetWorkspaceAgentStartupLogsAfter(_ context.Context, arg d
 	return logs, nil
 }
 
+func (q *fakeQuerier) GetWorkspaceAgentStartupLogsEOF(_ context.Context, agentID uuid.UUID) (bool, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	var lastLog database.WorkspaceAgentStartupLog
+	for _, log := range q.workspaceAgentLogs {
+		if log.AgentID != agentID {
+			continue
+		}
+		if log.ID > lastLog.ID {
+			lastLog = log
+		}
+	}
+	return lastLog.EOF, nil
+}
+
 func (q *fakeQuerier) GetWorkspaceAgentStats(_ context.Context, createdAfter time.Time) ([]database.GetWorkspaceAgentStatsRow, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
@@ -4013,7 +4029,7 @@ func (q *fakeQuerier) InsertWorkspaceAgentStartupLogs(_ context.Context, arg dat
 	defer q.mutex.Unlock()
 
 	logs := []database.WorkspaceAgentStartupLog{}
-	id := int64(1)
+	id := int64(0)
 	if len(q.workspaceAgentLogs) > 0 {
 		id = q.workspaceAgentLogs[len(q.workspaceAgentLogs)-1].ID
 	}
@@ -4026,6 +4042,7 @@ func (q *fakeQuerier) InsertWorkspaceAgentStartupLogs(_ context.Context, arg dat
 			CreatedAt: arg.CreatedAt[index],
 			Level:     arg.Level[index],
 			Output:    output,
+			EOF:       arg.EOF[index],
 		})
 		outputLength += int32(len(output))
 	}

coderd/database/dbmetrics/dbmetrics.go

@@ -738,6 +738,13 @@ func (m metricsStore) GetWorkspaceAgentStartupLogsAfter(ctx context.Context, arg
 	return logs, err
 }
 
+func (m metricsStore) GetWorkspaceAgentStartupLogsEOF(ctx context.Context, agentID uuid.UUID) (bool, error) {
+	start := time.Now()
+	r0, r1 := m.s.GetWorkspaceAgentStartupLogsEOF(ctx, agentID)
+	m.queryLatencies.WithLabelValues("GetWorkspaceAgentStartupLogsEOF").Observe(time.Since(start).Seconds())
+	return r0, r1
+}
+
 func (m metricsStore) GetWorkspaceAgentStats(ctx context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentStatsRow, error) {
 	start := time.Now()
 	stats, err := m.s.GetWorkspaceAgentStats(ctx, createdAt)

coderd/database/dbmock/dbmock.go

@@ -1453,6 +1453,21 @@ func (mr *MockStoreMockRecorder) GetWorkspaceAgentStartupLogsAfter(arg0, arg1 in
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentStartupLogsAfter", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentStartupLogsAfter), arg0, arg1)
 }
 
+// GetWorkspaceAgentStartupLogsEOF mocks base method.
+func (m *MockStore) GetWorkspaceAgentStartupLogsEOF(arg0 context.Context, arg1 uuid.UUID) (bool, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetWorkspaceAgentStartupLogsEOF", arg0, arg1)
+	ret0, _ := ret[0].(bool)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetWorkspaceAgentStartupLogsEOF indicates an expected call of GetWorkspaceAgentStartupLogsEOF.
+func (mr *MockStoreMockRecorder) GetWorkspaceAgentStartupLogsEOF(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentStartupLogsEOF", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentStartupLogsEOF), arg0, arg1)
+}
+
 // GetWorkspaceAgentStats mocks base method.
 func (m *MockStore) GetWorkspaceAgentStats(arg0 context.Context, arg1 time.Time) ([]database.GetWorkspaceAgentStatsRow, error) {
 	m.ctrl.T.Helper()

coderd/database/dump.sql

@@ -548,9 +548,12 @@ CREATE TABLE workspace_agent_startup_logs (
     created_at timestamp with time zone NOT NULL,
     output character varying(1024) NOT NULL,
     id bigint NOT NULL,
-    level log_level DEFAULT 'info'::log_level NOT NULL
+    level log_level DEFAULT 'info'::log_level NOT NULL,
+    eof boolean DEFAULT false NOT NULL
 );
 
+COMMENT ON COLUMN workspace_agent_startup_logs.eof IS 'End of file reached';
+
 CREATE SEQUENCE workspace_agent_startup_logs_id_seq
     START WITH 1
     INCREMENT BY 1

coderd/database/migrations/000127_add_startup_logs_eof.down.sql

@@ -0,0 +1 @@
+ALTER TABLE workspace_agent_startup_logs DROP COLUMN eof;

coderd/database/migrations/000127_add_startup_logs_eof.up.sql

@@ -0,0 +1,3 @@
+ALTER TABLE workspace_agent_startup_logs ADD COLUMN eof boolean NOT NULL DEFAULT false;
+
+COMMENT ON COLUMN workspace_agent_startup_logs.eof IS 'End of file reached';

coderd/database/models.go

@@ -1731,6 +1731,8 @@ type WorkspaceAgentStartupLog struct {
 	Output    string    `db:"output" json:"output"`
 	ID        int64     `db:"id" json:"id"`
 	Level     LogLevel  `db:"level" json:"level"`
+	// End of file reached
+	EOF bool `db:"eof" json:"eof"`
 }
 
 type WorkspaceAgentStat struct {

coderd/database/querier.go

@@ -127,6 +127,7 @@ type sqlcQuerier interface {
 	GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanceID string) (WorkspaceAgent, error)
 	GetWorkspaceAgentMetadata(ctx context.Context, workspaceAgentID uuid.UUID) ([]WorkspaceAgentMetadatum, error)
 	GetWorkspaceAgentStartupLogsAfter(ctx context.Context, arg GetWorkspaceAgentStartupLogsAfterParams) ([]WorkspaceAgentStartupLog, error)
+	GetWorkspaceAgentStartupLogsEOF(ctx context.Context, agentID uuid.UUID) (bool, error)
 	GetWorkspaceAgentStats(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentStatsRow, error)
 	GetWorkspaceAgentStatsAndLabels(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentStatsAndLabelsRow, error)
 	GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceAgent, error)

coderd/database/querier_test.go

@@ -114,6 +114,7 @@ func TestInsertWorkspaceAgentStartupLogs(t *testing.T) {
 		CreatedAt: []time.Time{database.Now()},
 		Output:    []string{"first"},
 		Level:     []database.LogLevel{database.LogLevelInfo},
+		EOF:       []bool{false},
 		// 1 MB is the max
 		OutputLength: 1 << 20,
 	})
@@ -125,6 +126,7 @@ func TestInsertWorkspaceAgentStartupLogs(t *testing.T) {
 		CreatedAt:    []time.Time{database.Now()},
 		Output:       []string{"second"},
 		Level:        []database.LogLevel{database.LogLevelInfo},
+		EOF:          []bool{false},
 		OutputLength: 1,
 	})
 	require.True(t, database.IsStartupLogsLimitError(err))

coderd/database/queries.sql.go

@@ -5438,7 +5438,7 @@ func (q *sqlQuerier) GetWorkspaceAgentMetadata(ctx context.Context, workspaceAge
 
 const getWorkspaceAgentStartupLogsAfter = `-- name: GetWorkspaceAgentStartupLogsAfter :many
 SELECT
-	agent_id, created_at, output, id, level
+	agent_id, created_at, output, id, level, eof
 FROM
 	workspace_agent_startup_logs
 WHERE
@@ -5468,6 +5468,7 @@ func (q *sqlQuerier) GetWorkspaceAgentStartupLogsAfter(ctx context.Context, arg
 			&i.Output,
 			&i.ID,
 			&i.Level,
+			&i.EOF,
 		); err != nil {
 			return nil, err
 		}
@@ -5482,6 +5483,26 @@ func (q *sqlQuerier) GetWorkspaceAgentStartupLogsAfter(ctx context.Context, arg
 	return items, nil
 }
 
+const getWorkspaceAgentStartupLogsEOF = `-- name: GetWorkspaceAgentStartupLogsEOF :one
+SELECT CASE WHEN EXISTS (
+	SELECT
+		agent_id, created_at, output, id, level, eof
+	FROM
+		workspace_agent_startup_logs
+	WHERE
+		agent_id = $1
+		AND eof = true
+	LIMIT 1
+) THEN TRUE ELSE FALSE END
+`
+
+func (q *sqlQuerier) GetWorkspaceAgentStartupLogsEOF(ctx context.Context, agentID uuid.UUID) (bool, error) {
+	row := q.db.QueryRowContext(ctx, getWorkspaceAgentStartupLogsEOF, agentID)
+	var column_1 bool
+	err := row.Scan(&column_1)
+	return column_1, err
+}
+
 const getWorkspaceAgentsByResourceIDs = `-- name: GetWorkspaceAgentsByResourceIDs :many
 SELECT
 	id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version, last_connected_replica_id, connection_timeout_seconds, troubleshooting_url, motd_file, lifecycle_state, startup_script_timeout_seconds, expanded_directory, shutdown_script, shutdown_script_timeout_seconds, startup_logs_length, startup_logs_overflowed, subsystem, startup_script_behavior
@@ -5833,16 +5854,17 @@ func (q *sqlQuerier) InsertWorkspaceAgentMetadata(ctx context.Context, arg Inser
 const insertWorkspaceAgentStartupLogs = `-- name: InsertWorkspaceAgentStartupLogs :many
 WITH new_length AS (
 	UPDATE workspace_agents SET
-	startup_logs_length = startup_logs_length + $5 WHERE workspace_agents.id = $1
+	startup_logs_length = startup_logs_length + $6 WHERE workspace_agents.id = $1
 )
 INSERT INTO
-		workspace_agent_startup_logs (agent_id, created_at, output, level)
+		workspace_agent_startup_logs (agent_id, created_at, output, level, eof)
 	SELECT
 		$1 :: uuid AS agent_id,
 		unnest($2 :: timestamptz [ ]) AS created_at,
 		unnest($3 :: VARCHAR(1024) [ ]) AS output,
-		unnest($4 :: log_level [ ]) AS level
-	RETURNING workspace_agent_startup_logs.agent_id, workspace_agent_startup_logs.created_at, workspace_agent_startup_logs.output, workspace_agent_startup_logs.id, workspace_agent_startup_logs.level
+		unnest($4 :: log_level [ ]) AS level,
+		unnest($5 :: boolean [ ]) AS eof
+	RETURNING workspace_agent_startup_logs.agent_id, workspace_agent_startup_logs.created_at, workspace_agent_startup_logs.output, workspace_agent_startup_logs.id, workspace_agent_startup_logs.level, workspace_agent_startup_logs.eof
 `
 
 type InsertWorkspaceAgentStartupLogsParams struct {
@@ -5850,6 +5872,7 @@ type InsertWorkspaceAgentStartupLogsParams struct {
 	CreatedAt    []time.Time `db:"created_at" json:"created_at"`
 	Output       []string    `db:"output" json:"output"`
 	Level        []LogLevel  `db:"level" json:"level"`
+	EOF          []bool      `db:"eof" json:"eof"`
 	OutputLength int32       `db:"output_length" json:"output_length"`
 }
 
@@ -5859,6 +5882,7 @@ func (q *sqlQuerier) InsertWorkspaceAgentStartupLogs(ctx context.Context, arg In
 		pq.Array(arg.CreatedAt),
 		pq.Array(arg.Output),
 		pq.Array(arg.Level),
+		pq.Array(arg.EOF),
 		arg.OutputLength,
 	)
 	if err != nil {
@@ -5874,6 +5898,7 @@ func (q *sqlQuerier) InsertWorkspaceAgentStartupLogs(ctx context.Context, arg In
 			&i.Output,
 			&i.ID,
 			&i.Level,
+			&i.EOF,
 		); err != nil {
 			return nil, err
 		}

coderd/database/queries/workspaceagents.sql

@@ -146,18 +146,31 @@ WHERE
 		id > @created_after
 	) ORDER BY id ASC;
 
+-- name: GetWorkspaceAgentStartupLogsEOF :one
+SELECT CASE WHEN EXISTS (
+	SELECT
+		*
+	FROM
+		workspace_agent_startup_logs
+	WHERE
+		agent_id = $1
+		AND eof = true
+	LIMIT 1
+) THEN TRUE ELSE FALSE END;
+
 -- name: InsertWorkspaceAgentStartupLogs :many
 WITH new_length AS (
 	UPDATE workspace_agents SET
 	startup_logs_length = startup_logs_length + @output_length WHERE workspace_agents.id = @agent_id
 )
 INSERT INTO
-		workspace_agent_startup_logs (agent_id, created_at, output, level)
+		workspace_agent_startup_logs (agent_id, created_at, output, level, eof)
 	SELECT
 		@agent_id :: uuid AS agent_id,
 		unnest(@created_at :: timestamptz [ ]) AS created_at,
 		unnest(@output :: VARCHAR(1024) [ ]) AS output,
-		unnest(@level :: log_level [ ]) AS level
+		unnest(@level :: log_level [ ]) AS level,
+		unnest(@eof :: boolean [ ]) AS eof
 	RETURNING workspace_agent_startup_logs.*;
 
 -- If an agent hasn't connected in the last 7 days, we purge it's logs.

coderd/database/sqlc.yaml

@@ -55,6 +55,7 @@ overrides:
       uuid: UUID
       failure_ttl: FailureTTL
       inactivity_ttl: InactivityTTL
+      eof: EOF
 
 sql:
   - schema: "./dump.sql"