coder · joobisb · Sep 21, 2024 · Oct 1, 2024
diff --git a/scaletest/harness/results.go b/scaletest/harness/results.go
@@ -27,14 +27,16 @@ type Results struct {
 
 // RunResult is the result of a single test run.
 type RunResult struct {
-	FullID     string           `json:"full_id"`
-	TestName   string           `json:"test_name"`
-	ID         string           `json:"id"`
-	Logs       string           `json:"logs"`
-	Error      error            `json:"error"`
-	StartedAt  time.Time        `json:"started_at"`
-	Duration   httpapi.Duration `json:"duration"`
-	DurationMS int64            `json:"duration_ms"`
+	FullID            string           `json:"full_id"`
+	TestName          string           `json:"test_name"`
+	ID                string           `json:"id"`
+	Logs              string           `json:"logs"`
+	Error             error            `json:"error"`
+	StartedAt         time.Time        `json:"started_at"`
+	Duration          httpapi.Duration `json:"duration"`
+	DurationMS        int64            `json:"duration_ms"`
+	TotalBytesRead    int64            `json:"total_bytes_read"`
+	TotalBytesWritten int64            `json:"total_bytes_written"`
 }
 
 // MarshalJSON implements json.Marhshaler for RunResult.
@@ -59,14 +61,16 @@ func (r *TestRun) Result() RunResult {
 	}
 
 	return RunResult{
-		FullID:     r.FullID(),
-		TestName:   r.testName,
-		ID:         r.id,
-		Logs:       r.logs.String(),
-		Error:      r.err,
-		StartedAt:  r.started,
-		Duration:   httpapi.Duration(r.duration),
-		DurationMS: r.duration.Milliseconds(),
+		FullID:            r.FullID(),
+		TestName:          r.testName,
+		ID:                r.id,
+		Logs:              r.logs.String(),
+		Error:             r.err,
+		StartedAt:         r.started,
+		Duration:          httpapi.Duration(r.duration),
+		DurationMS:        r.duration.Milliseconds(),
+		TotalBytesRead:    r.bytesRead,
+		TotalBytesWritten: r.bytesWritten,
 	}
 }
 

diff --git a/scaletest/harness/results_test.go b/scaletest/harness/results_test.go
@@ -36,34 +36,40 @@ func Test_Results(t *testing.T) {
 		TotalFail: 2,
 		Runs: map[string]harness.RunResult{
 			"test-0/0": {
-				FullID:     "test-0/0",
-				TestName:   "test-0",
-				ID:         "0",
-				Logs:       "test-0/0 log line 1\ntest-0/0 log line 2",
-				Error:      xerrors.New("test-0/0 error"),
-				StartedAt:  now,
-				Duration:   httpapi.Duration(time.Second),
-				DurationMS: 1000,
+				FullID:            "test-0/0",
+				TestName:          "test-0",
+				ID:                "0",
+				Logs:              "test-0/0 log line 1\ntest-0/0 log line 2",
+				Error:             xerrors.New("test-0/0 error"),
+				StartedAt:         now,
+				Duration:          httpapi.Duration(time.Second),
+				DurationMS:        1000,
+				TotalBytesRead:    1024,
+				TotalBytesWritten: 2048,
 			},
 			"test-0/1": {
-				FullID:     "test-0/1",
-				TestName:   "test-0",
-				ID:         "1",
-				Logs:       "test-0/1 log line 1\ntest-0/1 log line 2",
-				Error:      nil,
-				StartedAt:  now.Add(333 * time.Millisecond),
-				Duration:   httpapi.Duration(time.Second),
-				DurationMS: 1000,
+				FullID:            "test-0/1",
+				TestName:          "test-0",
+				ID:                "1",
+				Logs:              "test-0/1 log line 1\ntest-0/1 log line 2",
+				Error:             nil,
+				StartedAt:         now.Add(333 * time.Millisecond),
+				Duration:          httpapi.Duration(time.Second),
+				DurationMS:        1000,
+				TotalBytesRead:    512,
+				TotalBytesWritten: 1024,
 			},
 			"test-0/2": {
-				FullID:     "test-0/2",
-				TestName:   "test-0",
-				ID:         "2",
-				Logs:       "test-0/2 log line 1\ntest-0/2 log line 2",
-				Error:      testError{hidden: xerrors.New("test-0/2 error")},
-				StartedAt:  now.Add(666 * time.Millisecond),
-				Duration:   httpapi.Duration(time.Second),
-				DurationMS: 1000,
+				FullID:            "test-0/2",
+				TestName:          "test-0",
+				ID:                "2",
+				Logs:              "test-0/2 log line 1\ntest-0/2 log line 2",
+				Error:             testError{hidden: xerrors.New("test-0/2 error")},
+				StartedAt:         now.Add(666 * time.Millisecond),
+				Duration:          httpapi.Duration(time.Second),
+				DurationMS:        1000,
+				TotalBytesRead:    2048,
+				TotalBytesWritten: 4096,
 			},
 		},
 		Elapsed:   httpapi.Duration(time.Second),
@@ -109,6 +115,8 @@ Test results:
 			"started_at": "2023-10-05T12:03:56.395813665Z",
 			"duration": "1s",
 			"duration_ms": 1000,
+			"total_bytes_read": 1024,
+			"total_bytes_written": 2048,
 			"error": "test-0/0 error:\n    github.com/coder/coder/v2/scaletest/harness_test.Test_Results\n        [working_directory]/results_test.go:43"
 		},
 		"test-0/1": {
@@ -119,6 +127,8 @@ Test results:
 			"started_at": "2023-10-05T12:03:56.728813665Z",
 			"duration": "1s",
 			"duration_ms": 1000,
+			"total_bytes_read": 512,
+			"total_bytes_written": 1024,
 			"error": "\u003cnil\u003e"
 		},
 		"test-0/2": {
@@ -129,6 +139,8 @@ Test results:
 			"started_at": "2023-10-05T12:03:57.061813665Z",
 			"duration": "1s",
 			"duration_ms": 1000,
+			"total_bytes_read": 2048,
+			"total_bytes_written": 4096,
 			"error": "test-0/2 error"
 		}
 	}

diff --git a/scaletest/harness/run.go b/scaletest/harness/run.go
@@ -31,6 +31,13 @@ type Cleanable interface {
 	Cleanup(ctx context.Context, id string, logs io.Writer) error
 }
 
+// Collectable is an optional extension to Runnable that allows to get metrics from the runner
+type Collectable interface {
+	Runnable
+	// Gets the bytes transferred
+	GetBytesTransferred() (int64, int64)
+}
+
 // AddRun creates a new *TestRun with the given name, ID and Runnable, adds it
 // to the harness and returns it. Panics if the harness has been started, or a
 // test with the given run.FullID() is already registered.
@@ -66,11 +73,13 @@ type TestRun struct {
 	id       string
 	runner   Runnable
 
-	logs     *syncBuffer
-	done     chan struct{}
-	started  time.Time
-	duration time.Duration
-	err      error
+	logs         *syncBuffer
+	done         chan struct{}
+	started      time.Time
+	duration     time.Duration
+	err          error
+	bytesRead    int64
+	bytesWritten int64
 }
 
 func NewTestRun(testName string, id string, runner Runnable) *TestRun {
@@ -98,6 +107,11 @@ func (r *TestRun) Run(ctx context.Context) (err error) {
 	defer func() {
 		r.duration = time.Since(r.started)
 		r.err = err
+		c, ok := r.runner.(Collectable)
+		if !ok {
+			return
+		}
+		r.bytesRead, r.bytesWritten = c.GetBytesTransferred()
 	}()
 	defer func() {
 		e := recover()
@@ -107,6 +121,7 @@ func (r *TestRun) Run(ctx context.Context) (err error) {
 	}()
 
 	err = r.runner.Run(ctx, r.id, r.logs)
+
 	//nolint:revive // we use named returns because we mutate it in a defer
 	return
 }

diff --git a/scaletest/harness/run_test.go b/scaletest/harness/run_test.go
@@ -17,13 +17,24 @@ type testFns struct {
 	RunFn func(ctx context.Context, id string, logs io.Writer) error
 	// CleanupFn is optional if no cleanup is required.
 	CleanupFn func(ctx context.Context, id string, logs io.Writer) error
+	// CollectableFn is optional if byte transfer tracking is required.
+	CollectableFn func() (int64, int64)
 }
 
 // Run implements Runnable.
 func (fns testFns) Run(ctx context.Context, id string, logs io.Writer) error {
 	return fns.RunFn(ctx, id, logs)
 }
 
+// GetBytesTransferred implements Collectable.
+func (fns testFns) GetBytesTransferred() (bytesRead int64, bytesWritten int64) {
+	if fns.CollectableFn == nil {
+		return 0, 0
+	}
+
+	return fns.CollectableFn()
+}
+
 // Cleanup implements Cleanable.
 func (fns testFns) Cleanup(ctx context.Context, id string, logs io.Writer) error {
 	if fns.CleanupFn == nil {
@@ -40,9 +51,10 @@ func Test_TestRun(t *testing.T) {
 		t.Parallel()
 
 		var (
-			name, id      = "test", "1"
-			runCalled     int64
-			cleanupCalled int64
+			name, id          = "test", "1"
+			runCalled         int64
+			cleanupCalled     int64
+			collectableCalled int64
 
 			testFns = testFns{
 				RunFn: func(ctx context.Context, id string, logs io.Writer) error {
@@ -53,6 +65,10 @@ func Test_TestRun(t *testing.T) {
 					atomic.AddInt64(&cleanupCalled, 1)
 					return nil
 				},
+				CollectableFn: func() (int64, int64) {
+					atomic.AddInt64(&collectableCalled, 1)
+					return 0, 0
+				},
 			}
 		)
 
@@ -62,6 +78,7 @@ func Test_TestRun(t *testing.T) {
 		err := run.Run(context.Background())
 		require.NoError(t, err)
 		require.EqualValues(t, 1, atomic.LoadInt64(&runCalled))
+		require.EqualValues(t, 1, atomic.LoadInt64(&collectableCalled))
 
 		err = run.Cleanup(context.Background())
 		require.NoError(t, err)
@@ -105,6 +122,24 @@ func Test_TestRun(t *testing.T) {
 		})
 	})
 
+	t.Run("Collectable", func(t *testing.T) {
+		t.Parallel()
+
+		t.Run("NoFn", func(t *testing.T) {
+			t.Parallel()
+
+			run := harness.NewTestRun("test", "1", testFns{
+				RunFn: func(ctx context.Context, id string, logs io.Writer) error {
+					return nil
+				},
+				CollectableFn: nil,
+			})
+
+			err := run.Run(context.Background())
+			require.NoError(t, err)
+		})
+	})
+
 	t.Run("CatchesRunPanic", func(t *testing.T) {
 		t.Parallel()
 

diff --git a/scaletest/workspacetraffic/metrics.go b/scaletest/workspacetraffic/metrics.go
@@ -75,12 +75,14 @@ type ConnMetrics interface {
 	AddError(float64)
 	ObserveLatency(float64)
 	AddTotal(float64)
+	GetTotal() int64
 }
 
 type connMetrics struct {
 	addError       func(float64)
 	observeLatency func(float64)
 	addTotal       func(float64)
+	total          int64
 }
 
 func (c *connMetrics) AddError(f float64) {
@@ -92,5 +94,10 @@ func (c *connMetrics) ObserveLatency(f float64) {
 }
 
 func (c *connMetrics) AddTotal(f float64) {
+	c.total += int64(f)
 	c.addTotal(f)
 }
+
+func (c *connMetrics) GetTotal() int64 {
+	return c.total
+}
diff --git a/scaletest/workspacetraffic/run.go b/scaletest/workspacetraffic/run.go
@@ -204,6 +204,12 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
 	}
 }
 
+func (r *Runner) GetBytesTransferred() (bytesRead, bytesWritten int64) {
+	bytesRead = r.cfg.ReadMetrics.GetTotal()
+	bytesWritten = r.cfg.WriteMetrics.GetTotal()
+	return bytesRead, bytesWritten
+}
+
 // Cleanup does nothing, successfully.
 func (*Runner) Cleanup(context.Context, string, io.Writer) error {
 	return nil

diff --git a/scaletest/workspacetraffic/run_test.go b/scaletest/workspacetraffic/run_test.go
@@ -423,3 +423,7 @@ func (m *testMetrics) Latencies() []float64 {
 	defer m.Unlock()
 	return m.latencies
 }
+
+func (m *testMetrics) GetTotal() int64 {
+	return int64(m.total)
+}