feat: add load testing harness, coder loadtest command (#4853)

2026-09-24 15:04:27 +08:00 · 2022-11-02 18:30:00 +00:00
parent b1c400a7df
commit e7dd3f9378
23 changed files with 2641 additions and 6 deletions
@@ -0,0 +1,179 @@
+package cli
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"strconv"
+	"time"
+
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/cli/cliflag"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/loadtest/harness"
+)
+
+func loadtest() *cobra.Command {
+	var (
+		configPath string
+	)
+	cmd := &cobra.Command{
+		Use:   "loadtest --config <path>",
+		Short: "Load test the Coder API",
+		// TODO: documentation and a JSON scheme file
+		Long: "Perform load tests against the Coder server. The load tests " +
+			"configurable via a JSON file.",
+		Hidden: true,
+		Args:   cobra.ExactArgs(0),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if configPath == "" {
+				return xerrors.New("config is required")
+			}
+
+			var (
+				configReader io.ReadCloser
+			)
+			if configPath == "-" {
+				configReader = io.NopCloser(cmd.InOrStdin())
+			} else {
+				f, err := os.Open(configPath)
+				if err != nil {
+					return xerrors.Errorf("open config file %q: %w", configPath, err)
+				}
+				configReader = f
+			}
+
+			var config LoadTestConfig
+			err := json.NewDecoder(configReader).Decode(&config)
+			_ = configReader.Close()
+			if err != nil {
+				return xerrors.Errorf("read config file %q: %w", configPath, err)
+			}
+
+			err = config.Validate()
+			if err != nil {
+				return xerrors.Errorf("validate config: %w", err)
+			}
+
+			client, err := CreateClient(cmd)
+			if err != nil {
+				return err
+			}
+
+			me, err := client.User(cmd.Context(), codersdk.Me)
+			if err != nil {
+				return xerrors.Errorf("fetch current user: %w", err)
+			}
+
+			// Only owners can do loadtests. This isn't a very strong check but
+			// there's not much else we can do. Ratelimits are enforced for
+			// non-owners so hopefully that limits the damage if someone
+			// disables this check and runs it against a non-owner account.
+			ok := false
+			for _, role := range me.Roles {
+				if role.Name == "owner" {
+					ok = true
+					break
+				}
+			}
+			if !ok {
+				return xerrors.Errorf("Not logged in as site owner. Load testing is only available to site owners.")
+			}
+
+			// Disable ratelimits for future requests.
+			client.BypassRatelimits = true
+
+			// Prepare the test.
+			strategy := config.Strategy.ExecutionStrategy()
+			th := harness.NewTestHarness(strategy)
+
+			for i, t := range config.Tests {
+				name := fmt.Sprintf("%s-%d", t.Type, i)
+
+				for j := 0; j < t.Count; j++ {
+					id := strconv.Itoa(j)
+					runner, err := t.NewRunner(client)
+					if err != nil {
+						return xerrors.Errorf("create %q runner for %s/%s: %w", t.Type, name, id, err)
+					}
+
+					th.AddRun(name, id, runner)
+				}
+			}
+
+			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "Running load test...")
+
+			testCtx := cmd.Context()
+			if config.Timeout > 0 {
+				var cancel func()
+				testCtx, cancel = context.WithTimeout(testCtx, time.Duration(config.Timeout))
+				defer cancel()
+			}
+
+			// TODO: live progress output
+			start := time.Now()
+			err = th.Run(testCtx)
+			if err != nil {
+				return xerrors.Errorf("run test harness (harness failure, not a test failure): %w", err)
+			}
+			elapsed := time.Since(start)
+
+			// Print the results.
+			// TODO: better result printing
+			// TODO: move result printing to the loadtest package, add multiple
+			//       output formats (like HTML, JSON)
+			res := th.Results()
+			var totalDuration time.Duration
+			for _, run := range res.Runs {
+				totalDuration += run.Duration
+				if run.Error == nil {
+					continue
+				}
+
+				_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\n== FAIL: %s\n\n", run.FullID)
+				_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tError: %s\n\n", run.Error)
+
+				// Print log lines indented.
+				_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tLog:\n")
+				rd := bufio.NewReader(bytes.NewBuffer(run.Logs))
+				for {
+					line, err := rd.ReadBytes('\n')
+					if err == io.EOF {
+						break
+					}
+					if err != nil {
+						_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\n\tLOG PRINT ERROR: %+v\n", err)
+					}
+
+					_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\t\t%s", line)
+				}
+			}
+
+			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "\n\nTest results:")
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tPass:  %d\n", res.TotalPass)
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tFail:  %d\n", res.TotalFail)
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tTotal: %d\n", res.TotalRuns)
+			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "")
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tTotal duration: %s\n", elapsed)
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\tAvg. duration:  %s\n", totalDuration/time.Duration(res.TotalRuns))
+
+			// Cleanup.
+			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "\nCleaning up...")
+			err = th.Cleanup(cmd.Context())
+			if err != nil {
+				return xerrors.Errorf("cleanup tests: %w", err)
+			}
+
+			return nil
+		},
+	}
+
+	cliflag.StringVarP(cmd.Flags(), &configPath, "config", "", "CODER_LOADTEST_CONFIG_PATH", "", "Path to the load test configuration file, or - to read from stdin.")
+	return cmd
+}
@@ -0,0 +1,136 @@
+package cli_test
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/cli"
+	"github.com/coder/coder/cli/clitest"
+	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/coderd/httpapi"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/loadtest/placebo"
+	"github.com/coder/coder/loadtest/workspacebuild"
+	"github.com/coder/coder/pty/ptytest"
+	"github.com/coder/coder/testutil"
+)
+
+func TestLoadTest(t *testing.T) {
+	t.Parallel()
+
+	t.Run("PlaceboFromStdin", func(t *testing.T) {
+		t.Parallel()
+
+		client := coderdtest.New(t, nil)
+		_ = coderdtest.CreateFirstUser(t, client)
+
+		config := cli.LoadTestConfig{
+			Strategy: cli.LoadTestStrategy{
+				Type: cli.LoadTestStrategyTypeLinear,
+			},
+			Tests: []cli.LoadTest{
+				{
+					Type:  cli.LoadTestTypePlacebo,
+					Count: 10,
+					Placebo: &placebo.Config{
+						Sleep: httpapi.Duration(10 * time.Millisecond),
+					},
+				},
+			},
+			Timeout: httpapi.Duration(testutil.WaitShort),
+		}
+
+		configBytes, err := json.Marshal(config)
+		require.NoError(t, err)
+
+		cmd, root := clitest.New(t, "loadtest", "--config", "-")
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t)
+		cmd.SetIn(bytes.NewReader(configBytes))
+		cmd.SetOut(pty.Output())
+		cmd.SetErr(pty.Output())
+
+		ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancelFunc()
+
+		done := make(chan any)
+		go func() {
+			errC := cmd.ExecuteContext(ctx)
+			assert.NoError(t, errC)
+			close(done)
+		}()
+		pty.ExpectMatch("Test results:")
+		pty.ExpectMatch("Pass:  10")
+		cancelFunc()
+		<-done
+	})
+
+	t.Run("WorkspaceBuildFromFile", 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, nil)
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+
+		config := cli.LoadTestConfig{
+			Strategy: cli.LoadTestStrategy{
+				Type:             cli.LoadTestStrategyTypeConcurrent,
+				ConcurrencyLimit: 2,
+			},
+			Tests: []cli.LoadTest{
+				{
+					Type:  cli.LoadTestTypeWorkspaceBuild,
+					Count: 2,
+					WorkspaceBuild: &workspacebuild.Config{
+						OrganizationID: user.OrganizationID,
+						UserID:         user.UserID.String(),
+						Request: codersdk.CreateWorkspaceRequest{
+							TemplateID: template.ID,
+						},
+					},
+				},
+			},
+			Timeout: httpapi.Duration(testutil.WaitLong),
+		}
+
+		d := t.TempDir()
+		configPath := filepath.Join(d, "/config.loadtest.json")
+		f, err := os.Create(configPath)
+		require.NoError(t, err)
+		defer f.Close()
+		err = json.NewEncoder(f).Encode(config)
+		require.NoError(t, err)
+		_ = f.Close()
+
+		cmd, root := clitest.New(t, "loadtest", "--config", configPath)
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t)
+		cmd.SetIn(pty.Input())
+		cmd.SetOut(pty.Output())
+		cmd.SetErr(pty.Output())
+
+		ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancelFunc()
+
+		done := make(chan any)
+		go func() {
+			errC := cmd.ExecuteContext(ctx)
+			assert.NoError(t, errC)
+			close(done)
+		}()
+		pty.ExpectMatch("Test results:")
+		pty.ExpectMatch("Pass:  2")
+		<-done
+		cancelFunc()
+	})
+}
@@ -0,0 +1,181 @@
+package cli
+
+import (
+	"time"
+
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/coderd/httpapi"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/loadtest/harness"
+	"github.com/coder/coder/loadtest/placebo"
+	"github.com/coder/coder/loadtest/workspacebuild"
+)
+
+// LoadTestConfig is the overall configuration for a call to `coder loadtest`.
+type LoadTestConfig struct {
+	Strategy LoadTestStrategy `json:"strategy"`
+	Tests    []LoadTest       `json:"tests"`
+	// Timeout sets a timeout for the entire test run, to control the timeout
+	// for each individual run use strategy.timeout.
+	Timeout httpapi.Duration `json:"timeout"`
+}
+
+type LoadTestStrategyType string
+
+const (
+	LoadTestStrategyTypeLinear     LoadTestStrategyType = "linear"
+	LoadTestStrategyTypeConcurrent LoadTestStrategyType = "concurrent"
+)
+
+type LoadTestStrategy struct {
+	// Type is the type of load test strategy to use. Strategies determine how
+	// to run tests concurrently.
+	Type LoadTestStrategyType `json:"type"`
+
+	// ConcurrencyLimit is the maximum number of concurrent runs. This only
+	// applies if type == "concurrent". Negative values disable the concurrency
+	// limit and attempts to perform all runs concurrently. The default value is
+	// 100.
+	ConcurrencyLimit int `json:"concurrency_limit"`
+
+	// Shuffle determines whether or not to shuffle the test runs before
+	// executing them.
+	Shuffle bool `json:"shuffle"`
+	// Timeout is the maximum amount of time to run each test for. This is
+	// independent of the timeout specified in the test run. A timeout of 0
+	// disables the timeout.
+	Timeout httpapi.Duration `json:"timeout"`
+}
+
+func (s LoadTestStrategy) ExecutionStrategy() harness.ExecutionStrategy {
+	var strategy harness.ExecutionStrategy
+	switch s.Type {
+	case LoadTestStrategyTypeLinear:
+		strategy = harness.LinearExecutionStrategy{}
+	case LoadTestStrategyTypeConcurrent:
+		limit := s.ConcurrencyLimit
+		if limit < 0 {
+			return harness.ConcurrentExecutionStrategy{}
+		}
+		if limit == 0 {
+			limit = 100
+		}
+		strategy = harness.ParallelExecutionStrategy{
+			Limit: limit,
+		}
+	default:
+		panic("unreachable, unknown strategy type " + s.Type)
+	}
+
+	if s.Timeout > 0 {
+		strategy = harness.TimeoutExecutionStrategyWrapper{
+			Timeout: time.Duration(s.Timeout),
+			Inner:   strategy,
+		}
+	}
+	if s.Shuffle {
+		strategy = harness.ShuffleExecutionStrategyWrapper{
+			Inner: strategy,
+		}
+	}
+
+	return strategy
+}
+
+type LoadTestType string
+
+const (
+	LoadTestTypePlacebo        LoadTestType = "placebo"
+	LoadTestTypeWorkspaceBuild LoadTestType = "workspacebuild"
+)
+
+type LoadTest struct {
+	// Type is the type of load test to run.
+	Type LoadTestType `json:"type"`
+	// Count is the number of test runs to execute with this configuration. If
+	// the count is 0 or negative, defaults to 1.
+	Count int `json:"count"`
+
+	// Placebo must be set if type == "placebo".
+	Placebo *placebo.Config `json:"placebo,omitempty"`
+	// WorkspaceBuild must be set if type == "workspacebuild".
+	WorkspaceBuild *workspacebuild.Config `json:"workspacebuild,omitempty"`
+}
+
+func (t LoadTest) NewRunner(client *codersdk.Client) (harness.Runnable, error) {
+	switch t.Type {
+	case LoadTestTypePlacebo:
+		if t.Placebo == nil {
+			return nil, xerrors.New("placebo config must be set")
+		}
+
+		return placebo.NewRunner(*t.Placebo), nil
+	case LoadTestTypeWorkspaceBuild:
+		if t.WorkspaceBuild == nil {
+			return nil, xerrors.Errorf("workspacebuild config must be set")
+		}
+
+		return workspacebuild.NewRunner(client, *t.WorkspaceBuild), nil
+	default:
+		return nil, xerrors.Errorf("unknown test type %q", t.Type)
+	}
+}
+
+func (c *LoadTestConfig) Validate() error {
+	err := c.Strategy.Validate()
+	if err != nil {
+		return xerrors.Errorf("validate strategy: %w", err)
+	}
+
+	for i, test := range c.Tests {
+		err := test.Validate()
+		if err != nil {
+			return xerrors.Errorf("validate test %d: %w", i, err)
+		}
+	}
+
+	return nil
+}
+
+func (s *LoadTestStrategy) Validate() error {
+	switch s.Type {
+	case LoadTestStrategyTypeLinear:
+	case LoadTestStrategyTypeConcurrent:
+	default:
+		return xerrors.Errorf("invalid load test strategy type: %q", s.Type)
+	}
+
+	if s.Timeout < 0 {
+		return xerrors.Errorf("invalid load test strategy timeout: %q", s.Timeout)
+	}
+
+	return nil
+}
+
+func (t *LoadTest) Validate() error {
+	switch t.Type {
+	case LoadTestTypePlacebo:
+		if t.Placebo == nil {
+			return xerrors.Errorf("placebo test type must specify placebo")
+		}
+
+		err := t.Placebo.Validate()
+		if err != nil {
+			return xerrors.Errorf("validate placebo: %w", err)
+		}
+	case LoadTestTypeWorkspaceBuild:
+		if t.WorkspaceBuild == nil {
+			return xerrors.New("workspacebuild test type must specify workspacebuild")
+		}
+
+		err := t.WorkspaceBuild.Validate()
+		if err != nil {
+			return xerrors.Errorf("validate workspacebuild: %w", err)
+		}
+	default:
+		return xerrors.Errorf("invalid load test type: %q", t.Type)
+	}
+
+	return nil
+}
@@ -70,6 +70,7 @@ func init() {
 }

 func Core() []*cobra.Command {
+	// Please re-sort this list alphabetically if you change it!
 	return []*cobra.Command{
 		configSSH(),
 		create(),
@@ -77,26 +78,27 @@ func Core() []*cobra.Command {
 		dotfiles(),
 		gitssh(),
 		list(),
+		loadtest(),
 		login(),
 		logout(),
 		parameters(),
 		portForward(),
 		publickey(),
+		rename(),
 		resetPassword(),
 		schedules(),
 		show(),
-		ssh(),
 		speedtest(),
+		ssh(),
 		start(),
 		state(),
 		stop(),
-		rename(),
 		templates(),
+		tokens(),
 		update(),
 		users(),
 		versionCmd(),
 		workspaceAgent(),
-		tokens(),
 	}
 }