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

coderd/httpapi/json.go

@@ -0,0 +1,50 @@
+package httpapi
+
+import (
+	"encoding/json"
+	"time"
+
+	"golang.org/x/xerrors"
+)
+
+// Duration wraps time.Duration and provides better JSON marshaling and
+// unmarshaling. The default time.Duration marshals as an integer and only
+// accepts integers when unmarshaling, which is not very user friendly as users
+// cannot write durations like "1h30m".
+//
+// This type marshals as a string like "1h30m", and unmarshals from either a
+// string or an integer.
+type Duration time.Duration
+
+var _ json.Marshaler = Duration(0)
+var _ json.Unmarshaler = (*Duration)(nil)
+
+// MarshalJSON implements json.Marshaler.
+func (d Duration) MarshalJSON() ([]byte, error) {
+	return json.Marshal(time.Duration(d).String())
+}
+
+// UnmarshalJSON implements json.Unmarshaler.
+func (d *Duration) UnmarshalJSON(b []byte) error {
+	var v interface{}
+	err := json.Unmarshal(b, &v)
+	if err != nil {
+		return xerrors.Errorf("unmarshal JSON value: %w", err)
+	}
+
+	switch value := v.(type) {
+	case float64:
+		*d = Duration(time.Duration(value))
+		return nil
+	case string:
+		tmp, err := time.ParseDuration(value)
+		if err != nil {
+			return xerrors.Errorf("parse duration %q: %w", value, err)
+		}
+
+		*d = Duration(tmp)
+		return nil
+	}
+
+	return xerrors.New("invalid duration")
+}

coderd/httpapi/json_test.go

@@ -0,0 +1,168 @@
+package httpapi_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/coderd/httpapi"
+)
+
+func TestDuration(t *testing.T) {
+	t.Parallel()
+
+	t.Run("MarshalJSON", func(t *testing.T) {
+		t.Parallel()
+
+		cases := []struct {
+			value    time.Duration
+			expected string
+		}{
+			{
+				value:    0,
+				expected: "0s",
+			},
+			{
+				value:    1 * time.Millisecond,
+				expected: "1ms",
+			},
+			{
+				value:    1 * time.Second,
+				expected: "1s",
+			},
+			{
+				value:    1 * time.Minute,
+				expected: "1m0s",
+			},
+			{
+				value:    1 * time.Hour,
+				expected: "1h0m0s",
+			},
+			{
+				value:    1*time.Hour + 1*time.Minute + 1*time.Second + 1*time.Millisecond,
+				expected: "1h1m1.001s",
+			},
+		}
+
+		for _, c := range cases {
+			c := c
+
+			t.Run(c.expected, func(t *testing.T) {
+				t.Parallel()
+
+				d := httpapi.Duration(c.value)
+				b, err := d.MarshalJSON()
+				require.NoError(t, err)
+				require.Equal(t, `"`+c.expected+`"`, string(b))
+			})
+		}
+	})
+
+	t.Run("UnmarshalJSON", func(t *testing.T) {
+		t.Parallel()
+
+		cases := []struct {
+			value    string
+			expected time.Duration
+		}{
+			{
+				value:    "0ms",
+				expected: 0,
+			},
+			{
+				value:    "0s",
+				expected: 0,
+			},
+			{
+				value:    "1ms",
+				expected: 1 * time.Millisecond,
+			},
+			{
+				value:    "1s",
+				expected: 1 * time.Second,
+			},
+			{
+				value:    "1m",
+				expected: 1 * time.Minute,
+			},
+			{
+				value:    "1m0s",
+				expected: 1 * time.Minute,
+			},
+			{
+				value:    "1h",
+				expected: 1 * time.Hour,
+			},
+			{
+				value:    "1h0m0s",
+				expected: 1 * time.Hour,
+			},
+			{
+				value:    "1h1m1.001s",
+				expected: 1*time.Hour + 1*time.Minute + 1*time.Second + 1*time.Millisecond,
+			},
+			{
+				value:    "1h1m1s1ms",
+				expected: 1*time.Hour + 1*time.Minute + 1*time.Second + 1*time.Millisecond,
+			},
+		}
+
+		for _, c := range cases {
+			c := c
+
+			t.Run(c.value, func(t *testing.T) {
+				t.Parallel()
+
+				var d httpapi.Duration
+				err := d.UnmarshalJSON([]byte(`"` + c.value + `"`))
+				require.NoError(t, err)
+				require.Equal(t, c.expected, time.Duration(d))
+			})
+		}
+	})
+
+	t.Run("UnmarshalJSONInt", func(t *testing.T) {
+		t.Parallel()
+
+		var d httpapi.Duration
+		err := d.UnmarshalJSON([]byte("12345"))
+		require.NoError(t, err)
+		require.EqualValues(t, 12345, d)
+	})
+
+	t.Run("UnmarshalJSONErrors", func(t *testing.T) {
+		t.Parallel()
+
+		cases := []struct {
+			value       string
+			errContains string
+		}{
+			{
+				value:       "not valid json (no double quotes)",
+				errContains: "unmarshal JSON value",
+			},
+			{
+				value:       `"not valid duration"`,
+				errContains: "parse duration",
+			},
+			{
+				value:       "{}",
+				errContains: "invalid duration",
+			},
+		}
+
+		for _, c := range cases {
+			c := c
+
+			t.Run(c.value, func(t *testing.T) {
+				t.Parallel()
+
+				var d httpapi.Duration
+				err := d.UnmarshalJSON([]byte(c.value))
+				require.Error(t, err)
+				require.Contains(t, err.Error(), c.errContains)
+			})
+		}
+	})
+}

coderd/provisionerjobs.go

@@ -376,9 +376,20 @@ type provisionerJobLogsMessage struct {
 
 func (api *API) followLogs(jobID uuid.UUID) (<-chan database.ProvisionerJobLog, func(), error) {
 	logger := api.Logger.With(slog.F("job_id", jobID))
-	bufferedLogs := make(chan database.ProvisionerJobLog, 128)
-	closeSubscribe, err := api.Pubsub.Subscribe(provisionerJobLogsChannel(jobID),
+
+	var (
+		closed       = make(chan struct{})
+		bufferedLogs = make(chan database.ProvisionerJobLog, 128)
+	)
+	closeSubscribe, err := api.Pubsub.Subscribe(
+		provisionerJobLogsChannel(jobID),
 		func(ctx context.Context, message []byte) {
+			select {
+			case <-closed:
+				return
+			default:
+			}
+
 			jlMsg := provisionerJobLogsMessage{}
 			err := json.Unmarshal(message, &jlMsg)
 			if err != nil {
@@ -399,9 +410,11 @@ func (api *API) followLogs(jobID uuid.UUID) (<-chan database.ProvisionerJobLog,
 			}
 			if jlMsg.EndOfLogs {
 				logger.Debug(ctx, "got End of Logs")
+				close(closed)
 				close(bufferedLogs)
 			}
-		})
+		},
+	)
 	if err != nil {
 		return nil, nil, err
 	}