Refactor antithesis robustness main

Signed-off-by: Marek Siarkowicz <[email protected]>

Author: serathius

tests/antithesis/test-template/robustness/main.go

@@ -18,13 +18,12 @@ package main
 
 import (
 	"context"
-	"log"
+	"fmt"
 	"os"
 	"sync"
 	"time"
 
 	"github.com/antithesishq/antithesis-sdk-go/assert"
-	"github.com/antithesishq/antithesis-sdk-go/random"
 	"golang.org/x/time/rate"
 
 	"go.etcd.io/etcd/tests/v3/robustness/client"
@@ -43,45 +42,37 @@ var (
 		ClientCount:                    3,
 		MaxNonUniqueRequestConcurrency: 3,
 	}
-	IDProvider         = identity.NewIDProvider()
-	LeaseIDStorage     = identity.NewLeaseIDStorage()
-	ConcurrencyLimiter = traffic.NewConcurrencyLimiter(profile.MaxNonUniqueRequestConcurrency)
 )
 
-// Connect returns a client connection to an etcd node
-func Connect() *client.RecordingClient {
-	hosts := []string{"etcd0:2379", "etcd1:2379", "etcd2:2379"}
-	cli, err := client.NewRecordingClient(hosts, IDProvider, time.Now())
-	if err != nil {
-		log.Fatalf("Failed to connect to etcd: %v", err)
-		// Antithesis Assertion: client should always be able to connect to an etcd host
-		host := random.RandomChoice(hosts)
-		assert.Unreachable("Client failed to connect to an etcd host", map[string]any{"host": host, "error": err})
-		os.Exit(1)
-	}
-	return cli
+func main() {
+	ctx := context.Background()
+	baseTime := time.Now()
+	duration := time.Duration(robustnessrand.RandRange(5, 60) * int64(time.Second))
+	testRobustness(ctx, baseTime, duration)
 }
 
-func testRobustness() {
-	ctx := context.Background()
-	var wg sync.WaitGroup
-	var mux sync.Mutex
-	runfor := time.Duration(robustnessrand.RandRange(5, 60) * int64(time.Second))
+func testRobustness(ctx context.Context, baseTime time.Time, duration time.Duration) {
 	limiter := rate.NewLimiter(rate.Limit(profile.MaximalQPS), profile.BurstableQPS)
-	finish := wrap(time.After(runfor))
-	reports := []report.ClientReport{}
+	finish := closeAfter(ctx, duration)
+	ids := identity.NewIDProvider()
+	storage := identity.NewLeaseIDStorage()
+	concurrencyLimiter := traffic.NewConcurrencyLimiter(profile.MaxNonUniqueRequestConcurrency)
+	hosts := []string{"etcd0:2379", "etcd1:2379", "etcd2:2379"}
 
-	for range profile.ClientCount {
+	reports := []report.ClientReport{}
+	var mux sync.Mutex
+	var wg sync.WaitGroup
+	for i := 0; i < profile.ClientCount; i++ {
+		c := connect(hosts[i%len(hosts)], ids, baseTime)
 		wg.Add(1)
-		c := Connect()
 		go func(c *client.RecordingClient) {
 			defer wg.Done()
 			defer c.Close()
 
 			traffic.EtcdAntithesis.RunTrafficLoop(ctx, c, limiter,
-				IDProvider,
-				LeaseIDStorage,
-				ConcurrencyLimiter,
+				ids,
+				storage,
+				concurrencyLimiter,
 				finish,
 			)
 			mux.Lock()
@@ -90,21 +81,30 @@ func testRobustness() {
 		}(c)
 	}
 	wg.Wait()
-	assert.Reachable("Completion robustness traffic generation", nil)
+	fmt.Println("Completed robustness traffic generation")
+	assert.Reachable("Completed robustness traffic generation", nil)
 }
 
-// wrap converts a receive-only channel to receive-only struct{} channel
-func wrap[T any](from <-chan T) <-chan struct{} {
+func connect(endpoint string, ids identity.Provider, baseTime time.Time) *client.RecordingClient {
+	cli, err := client.NewRecordingClient([]string{endpoint}, ids, baseTime)
+	if err != nil {
+		// Antithesis Assertion: client should always be able to connect to an etcd host
+		assert.Unreachable("Client failed to connect to an etcd host", map[string]any{"host": endpoint, "error": err})
+		os.Exit(1)
+	}
+	return cli
+}
+
+func closeAfter(ctx context.Context, t time.Duration) <-chan struct{} {
 	out := make(chan struct{})
 	go func() {
 		for {
-			<-from
-			out <- struct{}{}
+			select {
+			case <-time.After(t):
+			case <-ctx.Done():
+			}
+			close(out)
 		}
 	}()
 	return out
 }
-
-func main() {
-	testRobustness()
-}