chore: allow multiple agent subsystems, add exectrace

coder · deansheather · Aug 9, 2023 · Aug 7, 2023 · Aug 8, 2023 · Aug 7, 2023
commit 6f18e5109a02e2a8edf76192ef4300b05926be89
diff --git a/agent/agent.go b/agent/agent.go
@@ -63,7 +63,7 @@ type Options struct {
 	IgnorePorts                  map[int]string
 	SSHMaxTimeout                time.Duration
 	TailnetListenPort            uint16
-	Subsystem                    codersdk.AgentSubsystem
+	Subsystems                   []codersdk.AgentSubsystem
 	Addresses                    []netip.Prefix
 	PrometheusRegistry           *prometheus.Registry
 	ReportMetadataInterval       time.Duration
@@ -144,7 +144,7 @@ func New(options Options) Agent {
 		reportMetadataInterval:       options.ReportMetadataInterval,
 		serviceBannerRefreshInterval: options.ServiceBannerRefreshInterval,
 		sshMaxTimeout:                options.SSHMaxTimeout,
-		subsystem:                    options.Subsystem,
+		subsystems:                   options.Subsystems,
 		addresses:                    options.Addresses,
 
 		prometheusRegistry: prometheusRegistry,
@@ -166,7 +166,7 @@ type agent struct {
 	// listing all listening ports. This is helpful to hide ports that
 	// are used by the agent, that the user does not care about.
 	ignorePorts map[int]string
-	subsystem   codersdk.AgentSubsystem
+	subsystems  []codersdk.AgentSubsystem
 
 	reconnectingPTYs       sync.Map
 	reconnectingPTYTimeout time.Duration
@@ -608,7 +608,7 @@ func (a *agent) run(ctx context.Context) error {
 	err = a.client.PostStartup(ctx, agentsdk.PostStartupRequest{
 		Version:           buildinfo.Version(),
 		ExpandedDirectory: manifest.Directory,
-		Subsystem:         a.subsystem,
+		Subsystems:        a.subsystems,
 	})
 	if err != nil {
 		return xerrors.Errorf("update workspace agent version: %w", err)

diff --git a/cli/agent.go b/cli/agent.go
@@ -12,6 +12,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -253,7 +254,19 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 			}
 
 			prometheusRegistry := prometheus.NewRegistry()
-			subsystem := inv.Environ.Get(agent.EnvAgentSubsystem)
+			subsystemsRaw := inv.Environ.Get(agent.EnvAgentSubsystem)
+			subsystems := []codersdk.AgentSubsystem{}
+			for _, s := range strings.Split(subsystemsRaw, ",") {
+				subsystem := codersdk.AgentSubsystem(strings.TrimSpace(s))
+				if subsystem == "" {
+					continue
+				}
+				if !subsystem.Valid() {
+					return xerrors.Errorf("invalid subsystem %q", subsystem)
+				}
+				subsystems = append(subsystems, subsystem)
+			}
+
 			agnt := agent.New(agent.Options{
 				Client:            client,
 				Logger:            logger,
@@ -275,7 +288,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 				},
 				IgnorePorts:   ignorePorts,
 				SSHMaxTimeout: sshMaxTimeout,
-				Subsystem:     codersdk.AgentSubsystem(subsystem),
+				Subsystems:    subsystems,
 
 				PrometheusRegistry: prometheusRegistry,
 			})

diff --git a/cli/agent_test.go b/cli/agent_test.go
@@ -2,6 +2,7 @@ package cli_test
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -264,8 +265,8 @@ func TestWorkspaceAgent(t *testing.T) {
 			"--agent-url", client.URL.String(),
 			"--log-dir", logDir,
 		)
-		// Set the subsystem for the agent.
-		inv.Environ.Set(agent.EnvAgentSubsystem, string(codersdk.AgentSubsystemEnvbox))
+		// Set the subsystems for the agent.
+		inv.Environ.Set(agent.EnvAgentSubsystem, fmt.Sprintf("%s,%s", codersdk.AgentSubsystemExectrace, codersdk.AgentSubsystemEnvbox))
 
 		pty := ptytest.New(t).Attach(inv)
 
@@ -275,6 +276,9 @@ func TestWorkspaceAgent(t *testing.T) {
 		resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
 		require.Len(t, resources, 1)
 		require.Len(t, resources[0].Agents, 1)
-		require.Equal(t, codersdk.AgentSubsystemEnvbox, resources[0].Agents[0].Subsystem)
+		require.Len(t, resources[0].Agents[0].Subsystems, 2)
+		// Sorted
+		require.Equal(t, codersdk.AgentSubsystemEnvbox, resources[0].Agents[0].Subsystems[0])
+		require.Equal(t, codersdk.AgentSubsystemExectrace, resources[0].Agents[0].Subsystems[1])
 	})
 }
diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -1064,8 +1064,10 @@ func (s *MethodTestSuite) TestWorkspace() {
 		res := dbgen.WorkspaceResource(s.T(), db, database.WorkspaceResource{JobID: build.JobID})
 		agt := dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{ResourceID: res.ID})
 		check.Args(database.UpdateWorkspaceAgentStartupByIDParams{
-			ID:        agt.ID,
-			Subsystem: database.WorkspaceAgentSubsystemNone,
+			ID: agt.ID,
+			Subsystems: []database.WorkspaceAgentSubsystem{
+				database.WorkspaceAgentSubsystemEnvbox,
+			},
 		}).Asserts(ws, rbac.ActionUpdate).Returns()
 	}))
 	s.Run("GetWorkspaceAgentLogsAfter", s.Subtest(func(db database.Store, check *expects) {

diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -5161,6 +5161,23 @@ func (q *FakeQuerier) UpdateWorkspaceAgentStartupByID(_ context.Context, arg dat
 		return err
 	}
 
+	if len(arg.Subsystems) > 0 {
+		seen := map[database.WorkspaceAgentSubsystem]struct{}{
+			arg.Subsystems[0]: {},
+		}
+		for i := 1; i < len(arg.Subsystems); i++ {
+			s := arg.Subsystems[i]
+			if _, ok := seen[s]; ok {
+				return xerrors.Errorf("duplicate subsystem %q", s)
+			}
+			seen[s] = struct{}{}
+
+			if arg.Subsystems[i-1] > arg.Subsystems[i] {
+				return xerrors.Errorf("subsystems not sorted: %q > %q", arg.Subsystems[i-1], arg.Subsystems[i])
+			}
+		}
+	}
+
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
 
@@ -5171,7 +5188,7 @@ func (q *FakeQuerier) UpdateWorkspaceAgentStartupByID(_ context.Context, arg dat
 
 		agent.Version = arg.Version
 		agent.ExpandedDirectory = arg.ExpandedDirectory
-		agent.Subsystem = arg.Subsystem
+		agent.Subsystems = arg.Subsystems
 		q.workspaceAgents[index] = agent
 		return nil
 	}

diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000148_agent_multiple_subsystems.down.sql b/coderd/database/migrations/000148_agent_multiple_subsystems.down.sql
@@ -0,0 +1,17 @@
+BEGIN;
+
+-- Bring back the subsytem column.
+ALTER TABLE workspace_agents ADD COLUMN subsystem workspace_agent_subsystem NOT NULL DEFAULT 'none';
+
+-- Update all existing workspace_agents to have subsystem = subsystems[0] unless
+-- subsystems is empty.
+UPDATE workspace_agents SET subsystem = subsystems[1] WHERE cardinality(subsystems) > 0;
+
+-- Drop the subsystems column from workspace_agents.
+ALTER TABLE workspace_agents DROP COLUMN subsystems;
+
+-- We cannot drop the "exectrace" value from the workspace_agent_subsystem type
+-- because you cannot drop values from an enum type.
+UPDATE workspace_agents SET subsystem = 'none' WHERE subsystem = 'exectrace';
+
+COMMIT;
diff --git a/coderd/database/migrations/000148_agent_multiple_subsystems.up.sql b/coderd/database/migrations/000148_agent_multiple_subsystems.up.sql
@@ -0,0 +1,21 @@
+BEGIN;
+
+-- Add "exectrace" to workspace_agent_subsystem type.
+ALTER TYPE workspace_agent_subsystem ADD VALUE 'exectrace';
+
+-- Create column subsystems in workspace_agents table, with default value being
+-- an empty array.
+ALTER TABLE workspace_agents ADD COLUMN subsystems workspace_agent_subsystem[] DEFAULT '{}';
+
+-- Add a constraint that the subsystems cannot contain the deprecated value
+-- 'none'.
+ALTER TABLE workspace_agents ADD CONSTRAINT subsystems_not_none CHECK (NOT ('none' = ANY (subsystems)));
+
+-- Update all existing workspace_agents to have subsystems = [subsystem] unless
+-- the subsystem is 'none'.
+UPDATE workspace_agents SET subsystems = ARRAY[subsystem] WHERE subsystem != 'none';
+
+-- Drop the subsystem column from workspace_agents.
+ALTER TABLE workspace_agents DROP COLUMN subsystem;
+
+COMMIT;
diff --git a/coderd/database/models.go b/coderd/database/models.go