diff --git a/enterprise/scaletest/agentfake/agent.go b/enterprise/scaletest/agentfake/agent.go new file mode 100644 index 0000000000..b03ebde8bd --- /dev/null +++ b/enterprise/scaletest/agentfake/agent.go @@ -0,0 +1,103 @@ +package agentfake + +import ( + "context" + "net/url" + "time" + + "golang.org/x/xerrors" + "google.golang.org/protobuf/types/known/timestamppb" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/agent/proto" + "github.com/coder/coder/v2/codersdk/agentsdk" +) + +const reconnectBackoff = 1 * time.Second + +// Agent is a single fake agent. It owns one workspace-agent auth token and one dRPC connection to coderd. +type Agent struct { + coderURL *url.URL + token string + logger slog.Logger + + cancel context.CancelFunc +} + +func NewAgent(coderURL *url.URL, token string, logger slog.Logger) *Agent { + return &Agent{ + coderURL: coderURL, + token: token, + logger: logger, + } +} + +// Run opens a dRPC websocket to coderd as the "agent" role and keeps it open until ctx is canceled or Close is called. +// On transient failures (e.g., coderd restart, brief auth churn while the workspace build is finalizing) Run reconnects +// with a small backoff. +// Returns nil when ctx is canceled or Close is called, and a non-nil error only if ctx returns a non-context error. +func (a *Agent) Run(ctx context.Context) error { + // Tie a.closed into ctx so a single select can wait on either. + runCtx, cancel := context.WithCancel(ctx) + a.cancel = cancel + defer a.cancel() + + client := agentsdk.New(a.coderURL, agentsdk.WithFixedToken(a.token)) + for { + if err := runCtx.Err(); err != nil { + return nil + } + err := a.connectAndServe(runCtx, client) + if err != nil && runCtx.Err() == nil { + a.logger.Warn(runCtx, "fake agent dRPC stream ended; reconnecting", + slog.Error(err)) + } + select { + case <-runCtx.Done(): + return nil + case <-time.After(reconnectBackoff): + } + } +} + +// connectAndServe opens one dRPC websocket, announces lifecycle = READY, then blocks until ctx is canceled or the +// connection is closed by either side. Returns the underlying error, if any. +func (a *Agent) connectAndServe(ctx context.Context, client *agentsdk.Client) error { + rpc, _, err := client.ConnectRPC28WithRole(ctx, "agent") + if err != nil { + return xerrors.Errorf("connect dRPC: %w", err) + } + conn := rpc.DRPCConn() + defer func() { + _ = conn.Close() + }() + + // Real agents transition to READY once their startup script finishes. Fakes have no startup script, so they're + // "ready" the moment the dRPC stream is open. We send this once per (re)connect because coderd's per-connection + // lifecycle state is reset each time. + // Failure here is logged but not treated as fatal: the connection itself is what flips Connected, and a transient + // failure to update lifecycle shouldn't tear the whole agent down. + if _, err := rpc.UpdateLifecycle(ctx, &proto.UpdateLifecycleRequest{ + Lifecycle: &proto.Lifecycle{ + State: proto.Lifecycle_READY, + ChangedAt: timestamppb.Now(), + }, + }); err != nil && ctx.Err() == nil { + a.logger.Warn(ctx, "failed to send lifecycle=READY", + slog.Error(err)) + } + + select { + case <-ctx.Done(): + return nil + case <-conn.Closed(): + return xerrors.New("dRPC connection closed by remote") + } +} + +// Close stops the agent. Safe to call multiple times. +func (a *Agent) Close() { + if a.cancel != nil { + a.cancel() + } +} diff --git a/enterprise/scaletest/agentfake/agent_test.go b/enterprise/scaletest/agentfake/agent_test.go new file mode 100644 index 0000000000..d01776f66d --- /dev/null +++ b/enterprise/scaletest/agentfake/agent_test.go @@ -0,0 +1,76 @@ +package agentfake_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "cdr.dev/slog/v3" + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/enterprise/scaletest/agentfake" + "github.com/coder/coder/v2/testutil" +) + +// Assert that our fake agent routine establishes the drpc connection and sets its lifecycle status to Ready. +func TestAgent_ConnectsAndReachesReady(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client, db := coderdtest.NewWithDatabase(t, nil) + user := coderdtest.CreateFirstUser(t, client) + + r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + }).WithAgent().Do() + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug) + a := agentfake.NewAgent(client.URL, r.AgentToken, logger) + t.Cleanup(func() { a.Close() }) + + runCtx, cancel := context.WithCancel(ctx) + t.Cleanup(cancel) + + runErr := make(chan error, 1) + go func() { + runErr <- a.Run(runCtx) + }() + + coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID). + WithContext(ctx). + Wait() + + require.Eventually(t, func() bool { + ws, err := client.Workspace(ctx, r.Workspace.ID) + if err != nil { + return false + } + for _, res := range ws.LatestBuild.Resources { + for _, agent := range res.Agents { + if agent.LifecycleState != codersdk.WorkspaceAgentLifecycleReady { + return false + } + } + } + return true + }, testutil.WaitLong, testutil.IntervalFast, + "agent never reached Lifecycle=ready in workspace %s", r.Workspace.ID) + + // Cancel Run and confirm a clean exit (nil error, not ctx error). + cancel() + select { + case err := <-runErr: + require.NoError(t, err, "Agent.Run returned unexpected error") + case <-ctx.Done(): + t.Fatalf("timed out waiting for Agent.Run to return: %v", ctx.Err()) + } + + // Close is idempotent and safe to call after Run returns. + a.Close() + a.Close() +} diff --git a/enterprise/scaletest/agentfake/manager.go b/enterprise/scaletest/agentfake/manager.go new file mode 100644 index 0000000000..d03e48307b --- /dev/null +++ b/enterprise/scaletest/agentfake/manager.go @@ -0,0 +1,219 @@ +package agentfake + +import ( + "context" + "errors" + "net/http" + "strconv" + "sync" + "time" + + "github.com/cenkalti/backoff/v4" + "github.com/google/uuid" + "golang.org/x/sync/errgroup" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/codersdk" +) + +const ( + enumeratePageSize = 100 + maxEnumerateRetries = 5 + initialEnumerateBackoff = 1 * time.Second + maxEnumerateRetryBackoff = 5 * time.Second +) + +// TokenInfo is a single workspace-agent auth token retrieved for a coder external agent, along with the identifying +// metadata needed to report the agent in metrics and logs. +type TokenInfo struct { + WorkspaceID uuid.UUID + WorkspaceName string + AgentID uuid.UUID + AgentName string + Token string +} + +// ManagerOptions configures a Manager. Authentication is supplied via the *codersdk.Client passed to NewManager rather +// than here, so the CLI / caller can construct the client with whatever session token (operator-issued, admin, +// template-admin) suits its deployment. +type ManagerOptions struct { + // Template restricts enumeration to workspaces of the given template name. Required. + Template string + // Owner restricts enumeration to workspaces owned by the given user. Optional; if empty, all owners are included. + Owner string +} + +// Manager supervises a set of fake Agents in one process. It enumerates the agents it owns from coderd at Run time +// (via coder_external_agent tokens on workspaces matching opts.Template), then opens a dRPC stream per agent and keeps +// them connected until ctx is canceled. +type Manager struct { + client *codersdk.Client + logger slog.Logger + opts ManagerOptions + + mu sync.Mutex + agents []*Agent +} + +// NewManager returns an Agent Manager. The provided client must already be authenticated with sufficient privilege +// to list workspaces by template and to call the enterprise-only WorkspaceExternalAgentCredentials endpoint +// (template-admin or higher; FeatureWorkspaceExternalAgent must be enabled). +func NewManager(client *codersdk.Client, logger slog.Logger, opts ManagerOptions) *Manager { + return &Manager{ + client: client, + logger: logger, + opts: opts, + } +} + +// Run enumerates the Manager's external agents from coderd, constructs one Agent per token, and runs the "fake agent" +// routines all until ctx is canceled or any Agent returns a non-context error. +// Enumeration is retried with exponential backoff for transient errors (network failures, 5xx, 429). +// Auth/permission/license/template-not-found errors (401, 403, 404) are treated as fatal. +// Run blocks until ctx is canceled, an Agent fails irrecoverably, or enumeration permanently fails. +func (m *Manager) Run(ctx context.Context) error { + if m.opts.Template == "" { + return xerrors.New("invalid manager options: Template is required") + } + + tokens, err := m.enumerateWithRetry(ctx) + if err != nil { + return xerrors.Errorf("enumerate external agents: %w", err) + } + + agents := make([]*Agent, 0, len(tokens)) + for i, ti := range tokens { + agents = append(agents, NewAgent(m.client.URL, ti.Token, + m.logger.Named("agent-"+strconv.Itoa(i)))) + } + m.mu.Lock() + m.agents = agents + m.mu.Unlock() + + eg, egCtx := errgroup.WithContext(ctx) + for _, a := range agents { + eg.Go(func() error { + return a.Run(egCtx) + }) + } + err = eg.Wait() + if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + return err + } + return nil +} + +// Close stops every Agent constructed during Run. Safe to call any +// number of times. +func (m *Manager) Close() { + for _, a := range m.agents { + a.Close() + } +} + +// enumerateWithRetry calls EnumerateExternalAgents with exponential backoff on transient failures. +// Fatal failures (auth, permission, missing template) exit immediately. +func (m *Manager) enumerateWithRetry(ctx context.Context) ([]TokenInfo, error) { + b := backoff.NewExponentialBackOff() + b.InitialInterval = initialEnumerateBackoff + b.MaxInterval = maxEnumerateRetryBackoff + bkoff := backoff.WithContext(backoff.WithMaxRetries(b, maxEnumerateRetries), ctx) + + var tokens []TokenInfo + // for attempt := 0; attempt <= maxEnumerateRetries; attempt++ { + err := backoff.Retry(func() error { + var retryErr error + tokens, retryErr = m.EnumerateExternalAgents(ctx) + if retryErr == nil { + return nil + } + if IsFatalEnumerationError(retryErr) { + m.logger.Warn(ctx, "enumeration failed, will retry", slog.Error(retryErr)) + return backoff.Permanent(retryErr) + } + return retryErr + }, bkoff) + if err != nil { + return nil, xerrors.Errorf("enumeration exhausted retries: %w", err) + } + return tokens, nil +} + +// EnumerateExternalAgents asks coderd for the list of workspaces matching the configured template, walks each +// workspace's latest build for agents on builds with HasExternalAgent=true, and returns the auth tokens for every +// external agent. Per-agent credential failures are logged and skipped; a non-nil error is returned only if the +// workspace listing itself fails. +func (m *Manager) EnumerateExternalAgents(ctx context.Context) ([]TokenInfo, error) { + var workspaces []codersdk.Workspace + filter := codersdk.WorkspaceFilter{ + Template: m.opts.Template, + Owner: m.opts.Owner, + Limit: enumeratePageSize, + } + for { + page, err := m.client.Workspaces(ctx, filter) + if err != nil { + return nil, xerrors.Errorf("list workspaces (offset=%d): %w", filter.Offset, err) + } + workspaces = append(workspaces, page.Workspaces...) + if len(page.Workspaces) < filter.Limit { + break + } + filter.Offset += len(page.Workspaces) + } + + tokens := make([]TokenInfo, 0, len(workspaces)) + for _, ws := range workspaces { + // The credentials endpoint requires WorkspaceBuild.HasExternalAgent=true (see + // enterprise/coderd/workspaceagents.go:48). Skip workspaces whose latest build + // doesn't carry the flag rather than 404 our way through every workspace in coderd. + if ws.LatestBuild.HasExternalAgent == nil || !*ws.LatestBuild.HasExternalAgent { + continue + } + for _, res := range ws.LatestBuild.Resources { + for _, agent := range res.Agents { + creds, err := m.client.WorkspaceExternalAgentCredentials(ctx, ws.ID, agent.Name) + if err != nil { + m.logger.Warn(ctx, "fetch external-agent credentials", + slog.F("workspace_id", ws.ID), + slog.F("workspace_name", ws.Name), + slog.F("agent_name", agent.Name), + slog.Error(err)) + continue + } + tokens = append(tokens, TokenInfo{ + WorkspaceID: ws.ID, + WorkspaceName: ws.Name, + AgentID: agent.ID, + AgentName: agent.Name, + Token: creds.AgentToken, + }) + } + } + } + return tokens, nil +} + +// IsFatalEnumerationError reports whether err from a coderd API call indicates an unrecoverable misconfiguration that +// retrying will not fix: missing/invalid session token, insufficient permissions, missing license feature, or a template +// that does not exist. +// All other errors (network blips, 429, 5xx) are treated as transient and can be retried. +func IsFatalEnumerationError(err error) bool { + if err == nil { + return false + } + sdkErr, ok := codersdk.AsError(err) + if !ok { + return false + } + + switch sdkErr.StatusCode() { + case http.StatusUnauthorized, + http.StatusForbidden, + http.StatusNotFound, + http.StatusBadRequest: + return true + } + return false +} diff --git a/enterprise/scaletest/agentfake/manager_test.go b/enterprise/scaletest/agentfake/manager_test.go new file mode 100644 index 0000000000..598729909f --- /dev/null +++ b/enterprise/scaletest/agentfake/manager_test.go @@ -0,0 +1,222 @@ +package agentfake_test + +import ( + "context" + "database/sql" + "sort" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "cdr.dev/slog/v3" + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" + "github.com/coder/coder/v2/enterprise/coderd/license" + "github.com/coder/coder/v2/enterprise/scaletest/agentfake" + sdkproto "github.com/coder/coder/v2/provisionersdk/proto" + "github.com/coder/coder/v2/testutil" +) + +// Asserts the TokenInfo shape (workspace IDs, agent names, tokens) returned by the enumeration loop. +func Test_Manager_EnumerateExternalAgents_returnsAllTokens(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client, db, user := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureWorkspaceExternalAgent: 1, + }, + }, + }) + + const numWorkspaces = 3 + first := buildExternalAgentWorkspace(t, db, user, uuid.Nil) + templateID := first.Workspace.TemplateID + want := []agentfake.TokenInfo{{ + WorkspaceID: first.Workspace.ID, + WorkspaceName: first.Workspace.Name, + AgentID: first.Agents[0].ID, + AgentName: first.Agents[0].Name, + Token: first.AgentToken, + }} + for i := 1; i < numWorkspaces; i++ { + r := buildExternalAgentWorkspace(t, db, user, templateID) + want = append(want, agentfake.TokenInfo{ + WorkspaceID: r.Workspace.ID, + WorkspaceName: r.Workspace.Name, + AgentID: r.Agents[0].ID, + AgentName: r.Agents[0].Name, + Token: r.AgentToken, + }) + } + + tmpl, err := client.Template(ctx, templateID) + require.NoError(t, err) + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug) + m := agentfake.NewManager(client, logger, agentfake.ManagerOptions{Template: tmpl.Name}) + + got, err := m.EnumerateExternalAgents(ctx) + require.NoError(t, err) + + // Order returned by coderd isn't guaranteed; sort both sides by WorkspaceID before comparing. + sortTokenInfosByWorkspaceID(want) + sortTokenInfosByWorkspaceID(got) + + require.Equal(t, len(want), len(got), + "expected one TokenInfo per external-agent workspace under the template") + for i := range want { + assert.Equal(t, want[i].WorkspaceID, got[i].WorkspaceID, "WorkspaceID for entry %d", i) + assert.Equal(t, want[i].AgentName, got[i].AgentName, "AgentName for entry %d", i) + assert.Equal(t, want[i].Token, got[i].Token, "Token for entry %d", i) + assert.NotEmpty(t, got[i].Token, "Token must be non-empty for entry %d", i) + } +} + +// Heavier-weight integration test for the agentfake harness: builds 5 external agents, sets up the client/Manager, +// and asserts that each of the agents the Manager sees via its enumeration function is properly connected and Ready. +func TestManager_FiveAgentsHeartbeat(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + + client, db, user := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureWorkspaceExternalAgent: 1, + }, + }, + }) + + const numAgents = 5 + first := buildExternalAgentWorkspace(t, db, user, uuid.Nil) + templateID := first.Workspace.TemplateID + workspaceIDs := []uuid.UUID{first.Workspace.ID} + for i := 1; i < numAgents; i++ { + r := buildExternalAgentWorkspace(t, db, user, templateID) + workspaceIDs = append(workspaceIDs, r.Workspace.ID) + } + + tmpl, err := client.Template(ctx, templateID) + require.NoError(t, err) + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug) + manager := agentfake.NewManager(client, logger, agentfake.ManagerOptions{ + Template: tmpl.Name, + }) + t.Cleanup(func() { manager.Close() }) + + managerCtx, cancelManager := context.WithCancel(ctx) + t.Cleanup(cancelManager) + + managerErr := make(chan error, 1) + go func() { + managerErr <- manager.Run(managerCtx) + }() + + // Each workspace's agent must reach Connected. Share the outer test ctx (testutil.WaitLong) across all five waiters + // so the total wait is bounded. + for _, wsID := range workspaceIDs { + coderdtest.NewWorkspaceAgentWaiter(t, client, wsID).WithContext(ctx).Wait() + } + + // Each workspace's agent must also reach Lifecycle=ready. The fake sends UpdateLifecycle(READY) once per dRPC + // connect; coderd persists that and exposes it on the agent. + for _, wsID := range workspaceIDs { + require.Eventually(t, func() bool { + ws, err := client.Workspace(ctx, wsID) + if err != nil { + return false + } + for _, res := range ws.LatestBuild.Resources { + for _, agent := range res.Agents { + if agent.LifecycleState != codersdk.WorkspaceAgentLifecycleReady { + return false + } + } + } + return true + }, testutil.WaitLong, testutil.IntervalFast, + "agent never reached Lifecycle=ready in workspace %s", wsID) + } + + // Cleanly stop the Manager and confirm it exits without a non-context error. + cancelManager() + select { + case err := <-managerErr: + if err != nil { + t.Fatalf("Manager.Run returned unexpected error: %v", err) + } + case <-ctx.Done(): + t.Fatalf("timed out waiting for Manager.Run to return: %v", ctx.Err()) + } +} + +// Asserts that an authentication failure during enumeration produces a fatal error, so the retry loop in +// enumerateWithRetry surfaces it immediately rather than hammering endpoints with credentials that will never work. +func Test_Manager_EnumerateExternalAgents_invalidTokenIsFatal(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client, db := coderdtest.NewWithDatabase(t, nil) + user := coderdtest.CreateFirstUser(t, client) + + r := buildExternalAgentWorkspace(t, db, user, uuid.Nil) + tmpl, err := client.Template(ctx, r.Workspace.TemplateID) + require.NoError(t, err) + + // Replace the client's session token with garbage to provoke a 401 from coderd's workspace-list endpoint. + // The Manager should surface that as a fatal error. + client.SetSessionToken("not-a-valid-session-token") + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug) + m := agentfake.NewManager(client, logger, agentfake.ManagerOptions{Template: tmpl.Name}) + + _, err = m.EnumerateExternalAgents(ctx) + require.Error(t, err, "expected enumeration to fail with an invalid session token") + require.True(t, agentfake.IsFatalEnumerationError(err), + "expected error to be classified as fatal so the harness exits and Kubernetes can restart it; got: %v", err) +} + +func sortTokenInfosByWorkspaceID(s []agentfake.TokenInfo) { + sort.Slice(s, func(i, j int) bool { + return s[i].WorkspaceID.String() < s[j].WorkspaceID.String() + }) +} + +// buildExternalAgentWorkspace creates one workspace with a coder_external_agent resource, an agent, and +// HasExternalAgent=true on the latest build. If templateID is uuid.Nil, dbfake mints a fresh template (and the caller +// can pass the returned Workspace.TemplateID into subsequent calls to share the template). +func buildExternalAgentWorkspace( + t *testing.T, + db database.Store, + user codersdk.CreateFirstUserResponse, + templateID uuid.UUID, +) dbfake.WorkspaceResponse { + t.Helper() + + ws := database.WorkspaceTable{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + } + if templateID != uuid.Nil { + ws.TemplateID = templateID + } + return dbfake.WorkspaceBuild(t, db, ws). + Seed(database.WorkspaceBuild{ + HasExternalAgent: sql.NullBool{Bool: true, Valid: true}, + }). + Resource(&sdkproto.Resource{ + Name: "external", + Type: "coder_external_agent", + }). + WithAgent(). + Do() +}