feat: add square wave latency summand

args.go

@@ -34,6 +34,12 @@ func parseArgs(args []string) (*lib.SpeedbumpCfg, error) {
 		sawPeriod = app.Flag("saw-period", "Period of the latency sawtooth wave.").
 				PlaceHolder("0").
 				Duration()
+		squareAmplitude = app.Flag("square-amplitude", "Amplitude of the latency square wave.").
+				PlaceHolder("0").
+				Duration()
+		suqarePeriod = app.Flag("square-period", "Period of the latency square wave.").
+				PlaceHolder("0").
+				Duration()
 		destAddr = app.Arg("destination", "TCP proxy destination in host:post format.").
 				Required().
 				String()
@@ -52,11 +58,13 @@ func parseArgs(args []string) (*lib.SpeedbumpCfg, error) {
 		BufferSize: int(*bufferSize),
 		QueueSize:  *queueSize,
 		Latency: &lib.LatencyCfg{
-			Base:          *latency,
-			SineAmplitude: *sineAmplitude,
-			SinePeriod:    *sinePeriod,
-			SawAmplitute:  *sawAmplitute,
-			SawPeriod:     *sawPeriod,
+			Base:            *latency,
+			SineAmplitude:   *sineAmplitude,
+			SinePeriod:    	 *sinePeriod,
+			SawAmplitute:  	 *sawAmplitute,
+			SawPeriod:     	 *sawPeriod,
+			SquareAmplitude:*squareAmplitude,
+			SquarePeriod: 	 *suqarePeriod,
 		},
 		LogLevel: *logLevel,
 	}

args_test.go

@@ -30,6 +30,8 @@ func TestParseArgsAll(t *testing.T) {
 			"--latency=100ms",
 			"--sine-amplitude=50ms",
 			"--sine-period=1m",
+			"--square-amplitude=123ms",
+			"--square-period=3m",
 			"host:777",
 		},
 	)
@@ -42,4 +44,6 @@ func TestParseArgsAll(t *testing.T) {
 	assert.Equal(t, time.Minute, cfg.Latency.SinePeriod)
 	assert.Equal(t, time.Duration(0), cfg.Latency.SawAmplitute)
 	assert.Equal(t, time.Duration(0), cfg.Latency.SawPeriod)
+	assert.Equal(t, time.Millisecond*123, cfg.Latency.SquareAmplitude)
+	assert.Equal(t, time.Minute*3, cfg.Latency.SquarePeriod)
 }

lib/latency_generator.go

@@ -9,11 +9,13 @@ type LatencyGenerator interface {
 }
 
 type LatencyCfg struct {
-	Base          time.Duration
-	SineAmplitude time.Duration
-	SinePeriod    time.Duration
-	SawAmplitute  time.Duration
-	SawPeriod     time.Duration
+	Base            time.Duration
+	SineAmplitude   time.Duration
+	SinePeriod      time.Duration
+	SawAmplitute    time.Duration
+	SawPeriod       time.Duration
+	SquareAmplitude time.Duration
+	SquarePeriod    time.Duration
 }
 
 type latencySummand interface {
@@ -39,6 +41,12 @@ func newSimpleLatencyGenerator(start time.Time, cfg *LatencyCfg) simpleLatencyGe
 			cfg.SawPeriod,
 		})
 	}
+	if cfg.SquareAmplitude > 0 && cfg.SquarePeriod > 0 {
+		summands = append(summands, squareLatencySummand{
+			cfg.SquareAmplitude,
+			cfg.SquarePeriod,
+		})
+	}
 	return simpleLatencyGenerator{
 		start:    start,
 		summands: summands,

lib/speedbump_test.go

@@ -140,7 +140,20 @@ func TestSpeedbumpWithEchoServer(t *testing.T) {
 	assert.Nil(t, err)
 
 	tcpAddr, _ := net.ResolveTCPAddr("tcp", "localhost:8000")
-	conn, _ := net.DialTCP("tcp", nil, tcpAddr)
+
+	var conn *net.TCPConn
+
+	// Wait for the speedbump instance to start listening
+	// since it is started in a separate goroutine, we don't know
+	// if it has already started listening by this point
+	for {
+		conn, err = net.DialTCP("tcp", nil, tcpAddr)
+		if err != nil {
+			time.Sleep(10 * time.Millisecond)
+		} else {
+			break
+		}
+	}
 
 	firstOpStart := time.Now()
 

lib/square.go

@@ -0,0 +1,17 @@
+package lib
+
+import (
+	"math"
+	"time"
+)
+
+type squareLatencySummand struct {
+	amplitude time.Duration
+	period    time.Duration
+}
+
+func (s squareLatencySummand) getLatency(elapsed time.Duration) time.Duration {
+	return time.Duration(
+		(2 * (2 * math.Floor((1 / float64 (s.period)) * float64(elapsed)) - math.Floor(2 * (1 / float64 (s.period)) * float64(elapsed))) + 1) * float64(s.amplitude),
+	)
+}

lib/square_test.go

@@ -0,0 +1,23 @@
+package lib
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSquareLatencySummand(t *testing.T) {
+	s := squareLatencySummand{
+		amplitude: time.Second * 2,
+		period:    time.Minute,
+	}
+
+	assert.Equal(t, s.getLatency(time.Duration(0)), time.Second * 2)
+	assert.Equal(t, s.getLatency(time.Second * 15), time.Second * 2)
+	assert.Equal(t, s.getLatency(time.Second * 30), time.Second * -2)
+	assert.Equal(t, s.getLatency(time.Second * 45), time.Second * -2)
+	assert.Equal(t, s.getLatency(time.Second * 60), time.Second * 2)
+	assert.Equal(t, s.getLatency(time.Second * 84), time.Second * 2)
+	assert.Equal(t, s.getLatency(time.Second * 90), time.Second * -2)
+}