feat: add subsystem to workspace agent

coder · sreya · May 18, 2023 · May 17, 2023 · May 17, 2023 · May 17, 2023
commit 90c85093613a5c86f6b775ab327f54e6633e0edc
diff --git a/agent/agent.go b/agent/agent.go
@@ -62,6 +62,7 @@ type Options struct {
 	IgnorePorts            map[int]string
 	SSHMaxTimeout          time.Duration
 	TailnetListenPort      uint16
+	Subsystem              agentsdk.AgentSubsystem
 }
 
 type Client interface {
@@ -119,6 +120,7 @@ func New(options Options) Agent {
 		ignorePorts:            options.IgnorePorts,
 		connStatsChan:          make(chan *agentsdk.Stats, 1),
 		sshMaxTimeout:          options.SSHMaxTimeout,
+		subsystem:              options.Subsystem,
 	}
 	a.init(ctx)
 	return a
@@ -136,6 +138,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   agentsdk.AgentSubsystem
 
 	reconnectingPTYs       sync.Map
 	reconnectingPTYTimeout time.Duration
@@ -1177,6 +1180,7 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
 		stats := &agentsdk.Stats{
 			ConnectionCount:    int64(len(networkStats)),
 			ConnectionsByProto: map[string]int64{},
+			Subsystem:          a.subsystem,
 		}
 		for conn, counts := range networkStats {
 			stats.ConnectionsByProto[conn.Proto.String()]++
@@ -1455,3 +1459,20 @@ func expandDirectory(dir string) (string, error) {
 	}
 	return dir, nil
 }
+
+// EnvAgentSubsystem is the environment variable used to denote the
+// specialized environment in which the agent is running
+// (e.g. envbox, envbuilder).
+const EnvAgentSubsytem = "CODER_AGENT_SUBSYSTEM"
+
+// SubsystemFromEnv returns the subsystem (if any) the agent
+// is running inside of.
+func SubsystemFromEnv() agentsdk.AgentSubsystem {
+	ss := os.Getenv(EnvAgentSubsytem)
+	switch agentsdk.AgentSubsystem(ss) {
+	case agentsdk.AgentSubsystemEnvbox:
+		return agentsdk.AgentSubsystemEnvbox
+	default:
+		return ""
+	}
+}
diff --git a/cli/agent.go b/cli/agent.go
@@ -218,6 +218,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 				},
 				IgnorePorts:   ignorePorts,
 				SSHMaxTimeout: sshMaxTimeout,
+				Subsystem:     agent.SubsystemFromEnv(),
 			})
 
 			debugSrvClose := ServeHandler(ctx, logger, agnt.HTTPDebug(), debugAddress, "debug")

diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000123_workspace_agent_subsystem.down.sql b/coderd/database/migrations/000123_workspace_agent_subsystem.down.sql
@@ -0,0 +1,4 @@
+BEGIN;
+ALTER TABLE workspace_agents DROP COLUMN subsystem;
+DROP TYPE workspace_agent_subsystem;
+COMMIT;
diff --git a/coderd/database/migrations/000123_workspace_agent_subsystem.up.sql b/coderd/database/migrations/000123_workspace_agent_subsystem.up.sql
@@ -0,0 +1,4 @@
+BEGIN;
+CREATE TYPE workspace_agent_subsystem AS ENUM ('envbuilder', 'envbox', 'none');
+ALTER TABLE workspace_agent_stats ADD COLUMN subsystem workspace_agent_subsystem NOT NULL default 'none';
+COMMIT;
diff --git a/coderd/database/models.go b/coderd/database/models.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaceagentstats.sql b/coderd/database/queries/workspaceagentstats.sql
@@ -17,10 +17,11 @@ INSERT INTO
 		session_count_jetbrains,
 		session_count_reconnecting_pty,
 		session_count_ssh,
-		connection_median_latency_ms
+		connection_median_latency_ms,
+		subsystem
 	)
 VALUES
-	($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17) RETURNING *;
+	($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18) RETURNING *;
 
 -- name: GetTemplateDAUs :many
 SELECT

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -1240,6 +1240,7 @@ func (api *API) workspaceAgentReportStats(rw http.ResponseWriter, r *http.Reques
 			SessionCountReconnectingPTY: req.SessionCountReconnectingPTY,
 			SessionCountSSH:             req.SessionCountSSH,
 			ConnectionMedianLatencyMS:   req.ConnectionMedianLatencyMS,
+			Subsystem:                   convertWorkspaceAgentSubsystem(req.Subsystem),
 		})
 		if err != nil {
 			return xerrors.Errorf("can't insert workspace agent stat: %w", err)
@@ -1983,3 +1984,12 @@ func convertWorkspaceAgentStartupLog(logEntry database.WorkspaceAgentStartupLog)
 		Level:     codersdk.LogLevel(logEntry.Level),
 	}
 }
+
+func convertWorkspaceAgentSubsystem(ss agentsdk.AgentSubsystem) database.WorkspaceAgentSubsystem {
+	switch ss {
+	case agentsdk.AgentSubsystemEnvbox:
+		return database.WorkspaceAgentSubsystemEnvbuilder
+	default:
+		return database.WorkspaceAgentSubsystemNone
+	}
+}
diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -453,6 +453,12 @@ func (c *Client) ReportStats(ctx context.Context, log slog.Logger, statsChan <-c
 	}), nil
 }
 
+type AgentSubsystem string
+
+const (
+	AgentSubsystemEnvbox AgentSubsystem = "envbox"
+)
+
 // Stats records the Agent's network connection statistics for use in
 // user-facing metrics and debugging.
 type Stats struct {
@@ -482,7 +488,8 @@ type Stats struct {
 	SessionCountReconnectingPTY int64 `json:"session_count_reconnecting_pty"`
 	// SessionCountSSH is the number of connections received by an agent
 	// that are normal, non-tagged SSH sessions.
-	SessionCountSSH int64 `json:"session_count_ssh"`
+	SessionCountSSH int64          `json:"session_count_ssh"`
+	Subsystem       AgentSubsystem `json:"subsystem"`
 
 	// Metrics collected by the agent
 	Metrics []AgentMetric `json:"metrics"`