workload: add upstream churn

cli/workload/upstreams.go

@@ -3,10 +3,12 @@ package workload
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"os"
 	"os/signal"
 	"strconv"
 	"syscall"
+	"time"
 
 	"github.com/andydunstall/piko/pkg/log"
 	"github.com/andydunstall/piko/workload/config"
@@ -87,7 +89,7 @@ func runUpstreams(conf *config.UpstreamsConfig, logger log.Logger) error {
 			logger,
 		)
 		g.Go(func() error {
-			return upstream.Run(ctx)
+			return runUpstream(ctx, upstream, conf)
 		})
 
 		nextEndpointID++
@@ -96,3 +98,39 @@ func runUpstreams(conf *config.UpstreamsConfig, logger log.Logger) error {
 
 	return g.Wait()
 }
+
+func runUpstream(
+	ctx context.Context,
+	upstream *upstream.Upstream,
+	conf *config.UpstreamsConfig,
+) error {
+	if conf.Churn.Interval == 0 {
+		return upstream.Run(ctx)
+	}
+
+	for {
+		multipler := rand.Float64()
+
+		churnInterval := time.Duration(float64(conf.Churn.Interval) * multipler)
+		upstreamCtx, cancel := context.WithTimeout(ctx, churnInterval)
+		defer cancel()
+
+		if err := upstream.Run(upstreamCtx); err != nil {
+			if upstreamCtx.Err() == nil {
+				return err
+			}
+		}
+
+		if ctx.Err() != nil {
+			return nil
+		}
+
+		if conf.Churn.Delay != 0 {
+			select {
+			case <-ctx.Done():
+				return nil
+			case <-time.After(conf.Churn.Delay):
+			}
+		}
+	}
+}

workload/config/upstreams.go

@@ -3,6 +3,7 @@ package config
 import (
 	"fmt"
 	"net/url"
+	"time"
 
 	"github.com/andydunstall/piko/pkg/log"
 	"github.com/spf13/pflag"
@@ -23,13 +24,42 @@ func (c *ServerConfig) Validate() error {
 	return nil
 }
 
+type ChurnConfig struct {
+	// Interval specifies how often each upstream should 'churn'
+	// (disconnect and reconnect).
+	Interval time.Duration `json:"interval" yaml:"interval"`
+
+	// Delay is the duration to wait before reconnecting when churning.
+	Delay time.Duration `json:"delay" yaml:"delay"`
+}
+
+func (c *ChurnConfig) RegisterFlags(fs *pflag.FlagSet) {
+	fs.DurationVar(
+		&c.Interval,
+		"churn.interval",
+		0,
+		`
+How often each upstream should 'churn' (disconnect and reconnect).`,
+	)
+
+	fs.DurationVar(
+		&c.Delay,
+		"churn.delay",
+		0,
+		`
+Duration to wait before reconnecting when an upstream churns.`,
+	)
+}
+
 type UpstreamsConfig struct {
 	// Upstreams is the number of upstream servers to register.
 	Upstreams int `json:"upstreams" yaml:"upstreams"`
 
 	// Endpoints is the number of endpoint IDs to register.
 	Endpoints int `json:"endpoints" yaml:"endpoints"`
 
+	Churn ChurnConfig `json:"churn" yaml:"churn"`
+
 	Server ServerConfig `json:"server" yaml:"server"`
 
 	Log log.Config `json:"log" yaml:"log"`
@@ -81,6 +111,8 @@ servers per endpoint.
 Therefore 'endpoints' must be greater than or equal to 'upstreams'.`,
 	)
 
+	c.Churn.RegisterFlags(fs)
+
 	fs.StringVar(
 		&c.Server.URL,
 		"server.url",

workload/upstream/upstream.go

@@ -34,11 +34,14 @@ func (u *Upstream) Run(ctx context.Context) error {
 	}))
 	defer server.Close()
 
-	client := client.New(client.WithUpstreamURL(u.serverURL))
+	client := client.New(
+		client.WithUpstreamURL(u.serverURL),
+		client.WithLogger(u.logger),
+	)
 
-	ln, err := client.Listen(context.Background(), u.endpointID)
+	ln, err := client.Listen(ctx, u.endpointID)
 	if err != nil {
-		return fmt.Errorf("listen: %s: %w", ln.EndpointID(), err)
+		return fmt.Errorf("listen: %w", err)
 	}
 	defer ln.Close()