Use robustness traffic generation

Signed-off-by: Nont <[email protected]>

Author: Nont

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

@@ -20,21 +20,38 @@ import (
 	"context"
 	"log"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/antithesishq/antithesis-sdk-go/assert"
 	"github.com/antithesishq/antithesis-sdk-go/random"
+	"golang.org/x/time/rate"
 
-	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.etcd.io/etcd/tests/v3/robustness/client"
 	"go.etcd.io/etcd/tests/v3/robustness/identity"
+	robustnessrand "go.etcd.io/etcd/tests/v3/robustness/random"
+	"go.etcd.io/etcd/tests/v3/robustness/report"
+	"go.etcd.io/etcd/tests/v3/robustness/traffic"
 )
 
-func Connect() *client.RecordingClient {
-	// This function returns a client connection to an etcd node
+var (
+	// Please keep the sum of weights equal 100.
+	profile = traffic.Profile{
+		MinimalQPS:                     100,
+		MaximalQPS:                     1000,
+		BurstableQPS:                   1000,
+		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, identity.NewIDProvider(), time.Now())
+	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
@@ -45,77 +62,49 @@ func Connect() *client.RecordingClient {
 	return cli
 }
 
-func DeleteKeys() {
-	// This function will:
-	// 1. Get all keys
-	// 2. Select half of the keys received
-	// 3. Attempt to delete the keys selected
-	// 4. Check that the keys were deleted
-
+func testRobustness() {
 	ctx := context.Background()
-
-	// Connect to an etcd node
-	cli := Connect()
-
-	// Get all keys
-	resp, err := cli.Get(ctx, "", clientv3.WithPrefix())
-
-	// Antithesis Assertion: sometimes get with prefix requests are successful. A failed request is OK since we expect them to happen.
-	assert.Sometimes(err == nil, "Client can make successful get all requests", map[string]any{"error": err})
-	cli.Close()
-
-	if err != nil {
-		log.Printf("Client failed to get all keys: %v", err)
-		os.Exit(0)
+	var wg sync.WaitGroup
+	var mux sync.Mutex
+	runfor := time.Duration(robustnessrand.RandRange(5, 60) * int64(time.Second))
+	limiter := rate.NewLimiter(rate.Limit(profile.MaximalQPS), profile.BurstableQPS)
+	finish := wrap(time.After(runfor))
+	reports := []report.ClientReport{}
+
+	for range profile.ClientCount {
+		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,
+				finish,
+			)
+			mux.Lock()
+			reports = append(reports, c.Report())
+			mux.Unlock()
+		}(c)
 	}
+	wg.Wait()
+	assert.Reachable("Completion robustness traffic generation", nil)
+}
 
-	// Choose half of the keys
-	var keys []string
-	for _, k := range resp.Kvs {
-		keys = append(keys, string(k.Key))
-	}
-	half := len(keys) / 2
-	halfKeys := keys[:half]
-
-	// Connect to a new etcd node
-	cli = Connect()
-
-	// Delete half of the keys chosen
-	var deletedKeys []string
-	for _, k := range halfKeys {
-		_, err := cli.Delete(ctx, k)
-		// Antithesis Assertion: sometimes delete requests are successful. A failed request is OK since we expect them to happen.
-		assert.Sometimes(err == nil, "Client can make successful delete requests", map[string]any{"error": err})
-		if err != nil {
-			log.Printf("Failed to delete key %s: %v", k, err)
-		} else {
-			log.Printf("Successfully deleted key %v", k)
-			deletedKeys = append(deletedKeys, k)
-		}
-	}
-	cli.Close()
-
-	// Connect to a new etcd node
-	cli = Connect()
-
-	// Check to see if those keys were deleted / exist
-	for _, k := range deletedKeys {
-		resp, err := cli.Get(ctx, k)
-		// Antithesis Assertion: sometimes get requests are successful. A failed request is OK since we expect them to happen.
-		assert.Sometimes(err == nil, "Client can make successful get requests", map[string]any{"error": err})
-		if err != nil {
-			log.Printf("Client failed to get key %s: %v", k, err)
-			continue
+// wrap converts a receive-only channel to receive-only struct{} channel
+func wrap[T any](from <-chan T) <-chan struct{} {
+	out := make(chan struct{})
+	go func() {
+		for {
+			<-from
+			out <- struct{}{}
 		}
-		// Antithesis Assertion: if we deleted a key, we should not get a value
-		assert.Always(resp.Count == 0, "Key was deleted correctly", map[string]any{"key": k})
-	}
-	cli.Close()
-
-	assert.Reachable("Completion of a key deleting check", nil)
-	log.Printf("Completion of a key deleting check")
+	}()
+	return out
 }
 
 func main() {
-	DeleteKeys()
+	testRobustness()
 }

tests/robustness/traffic/etcd.go

@@ -75,6 +75,14 @@ var (
 			{Choice: Delete, Weight: 50},
 		},
 	}
+	EtcdAntithesis Traffic = etcdTraffic{
+		keyCount: 10,
+		// Please keep the sum of weights equal 100.
+		requests: []random.ChoiceWeight[etcdRequestType]{
+			{Choice: Get, Weight: 50},
+			{Choice: Put, Weight: 50},
+		},
+	}
 )
 
 type etcdTraffic struct {