refactor bytes per second to bytes per tick and tick interval

coder · johnstcn · May 5, 2023 · Apr 18, 2023 · Apr 27, 2023 · Apr 27, 2023
commit 9dc28a2ba6f33b55a839a4f7b7adf4f9f1d09e5f
diff --git a/cli/scaletest.go b/cli/scaletest.go
@@ -897,7 +897,8 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd {
 func (r *RootCmd) scaletestTrafficGen() *clibase.Cmd {
 	var (
 		duration        time.Duration
-		bps             int64
+		tickInterval    time.Duration
+		bytesPerTick    int64
 		client          = &codersdk.Client{}
 		tracingFlags    = &scaletestTracingFlags{}
 		strategy        = &scaletestStrategyFlags{}
@@ -975,10 +976,10 @@ func (r *RootCmd) scaletestTrafficGen() *clibase.Cmd {
 
 				// Setup our workspace agent connection.
 				config := trafficgen.Config{
-					AgentID:        agentID,
-					BytesPerSecond: bps,
-					Duration:       duration,
-					TicksPerSecond: 10,
+					AgentID:      agentID,
+					BytesPerTick: bytesPerTick,
+					Duration:     duration,
+					TickInterval: tickInterval,
 				}
 
 				if err := config.Validate(); err != nil {
@@ -1029,11 +1030,18 @@ func (r *RootCmd) scaletestTrafficGen() *clibase.Cmd {
 			Value:       clibase.DurationOf(&duration),
 		},
 		{
-			Flag:        "bps",
-			Env:         "CODER_SCALETEST_TRAFFICGEN_BPS",
+			Flag:        "bytes-per-tick",
+			Env:         "CODER_SCALETEST_TRAFFICGEN_BYTES_PER_TICK",
 			Default:     "1024",
-			Description: "How much traffic to generate in bytes per second.",
-			Value:       clibase.Int64Of(&bps),
+			Description: "How much traffic to generate per tick.",
+			Value:       clibase.Int64Of(&bytesPerTick),
+		},
+		{
+			Flag:        "tick-interval",
+			Env:         "CODER_SCALETEST_TRAFFICGEN_TICK_INTERVAL",
+			Default:     "100ms",
+			Description: "How often to send traffic.",
+			Value:       clibase.DurationOf(&tickInterval),
 		},
 	}
 

diff --git a/cli/scaletest_test.go b/cli/scaletest_test.go
@@ -261,7 +261,8 @@ func TestScaleTestTrafficGen(t *testing.T) {
 
 	inv, root := clitest.New(t, "scaletest", "trafficgen", ws.Name,
 		"--duration", "1s",
-		"--bps", "100",
+		"--bytes-per-tick", "1024",
+		"--tick-interval", "100ms",
 	)
 	clitest.SetupConfig(t, client, root)
 	var stdout, stderr bytes.Buffer

diff --git a/scaletest/trafficgen/config.go b/scaletest/trafficgen/config.go
@@ -10,32 +10,33 @@ import (
 type Config struct {
 	// AgentID is the workspace agent ID to which to connect.
 	AgentID uuid.UUID `json:"agent_id"`
-	// BytesPerSecond is the number of bytes to send to the agent.
 
-	BytesPerSecond int64 `json:"bytes_per_second"`
+	// BytesPerTick is the number of bytes to send to the agent per tick.
+	BytesPerTick int64 `json:"bytes_per_tick"`
 
 	// Duration is the total duration for which to send traffic to the agent.
 	Duration time.Duration `json:"duration"`
 
-	// TicksPerSecond specifies how many times per second we send traffic.
-	TicksPerSecond int64 `json:"ticks_per_second"`
+	// TicksInterval specifies how many times per second we send traffic.
+	TickInterval time.Duration `json:"tick_interval"`
 }
 
 func (c Config) Validate() error {
 	if c.AgentID == uuid.Nil {
 		return xerrors.Errorf("validate agent_id: must not be nil")
 	}
 
-	if c.BytesPerSecond <= 0 {
-		return xerrors.Errorf("validate bytes_per_second: must be greater than zero")
+	if c.BytesPerTick <= 0 {
+		return xerrors.Errorf("validate bytes_per_tick: must be greater than zero")
 	}
 
 	if c.Duration <= 0 {
 		return xerrors.Errorf("validate duration: must be greater than zero")
 	}
 
-	if c.TicksPerSecond <= 0 {
-		return xerrors.Errorf("validate ticks_per_second: must be greater than zero")
+	if c.TickInterval <= 0 {
+		return xerrors.Errorf("validate tick_interval: must be greater than zero")
 	}
+
 	return nil
 }
diff --git a/scaletest/trafficgen/run.go b/scaletest/trafficgen/run.go
@@ -52,8 +52,8 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) error {
 		reconnect           = uuid.New()
 		height       uint16 = 25
 		width        uint16 = 80
-		tickInterval        = time.Second / time.Duration(r.cfg.TicksPerSecond)
-		bytesPerTick        = r.cfg.BytesPerSecond / r.cfg.TicksPerSecond
+		tickInterval        = r.cfg.TickInterval
+		bytesPerTick        = r.cfg.BytesPerTick
 	)
 
 	logger.Info(ctx, "config",
@@ -149,7 +149,10 @@ func (*Runner) Cleanup(context.Context, string) error {
 // drain drains from src until it returns io.EOF or ctx times out.
 func drain(src io.Reader) error {
 	if _, err := io.Copy(io.Discard, src); err != nil {
-		if xerrors.Is(err, context.DeadlineExceeded) || xerrors.Is(err, websocket.CloseError{}) {
+		if xerrors.Is(err, context.DeadlineExceeded) {
+			return nil
+		}
+		if xerrors.As(err, &websocket.CloseError{}) {
 			return nil
 		}
 		return err
@@ -166,7 +169,10 @@ func writeRandomData(dst io.Writer, size int64, tick <-chan time.Time) error {
 		payload := "#" + mustRandStr(size-1)
 		ptyReq.Data = payload
 		if err := enc.Encode(ptyReq); err != nil {
-			if xerrors.Is(err, context.DeadlineExceeded) || xerrors.Is(err, websocket.CloseError{}) {
+			if xerrors.Is(err, context.DeadlineExceeded) {
+				return nil
+			}
+			if xerrors.As(err, &websocket.CloseError{}) {
 				return nil
 			}
 			return err
@@ -213,8 +219,11 @@ func (w *countReadWriter) BytesWritten() int64 {
 	return w.bytesWritten.Load()
 }
 
-func mustRandStr(len int64) string {
-	randStr, err := cryptorand.String(int(len))
+func mustRandStr(l int64) string {
+	if l < 1 {
+		l = 1
+	}
+	randStr, err := cryptorand.String(int(l))
 	if err != nil {
 		panic(err)
 	}