Use robustness traffic generation

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

Author: Nont

tests/antithesis/test-template/go-delete-keys/serial_driver_delete_keys.go

@@ -20,14 +20,18 @@ 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 {
@@ -45,77 +49,64 @@ 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
-
-	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)
-	}
-
-	// 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]
+var etcdTraffic traffic.Traffic = traffic.EtcdTraffic{
+	KeyCount: 10,
+	// Please keep the sum of weights equal 100.
+	Requests: []robustnessrand.ChoiceWeight[traffic.EtcdRequestType]{
+		{Choice: traffic.Get, Weight: 50},
+		{Choice: traffic.Put, Weight: 50},
+	},
+}
 
-	// Connect to a new etcd node
-	cli = Connect()
+var profile = traffic.Profile{
+	MinimalQPS:                     100,
+	MaximalQPS:                     1000,
+	BurstableQPS:                   1000,
+	ClientCount:                    3,
+	MaxNonUniqueRequestConcurrency: 3,
+}
 
-	// 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)
+// 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{}{}
 		}
-	}
-	cli.Close()
-
-	// Connect to a new etcd node
-	cli = Connect()
+	}()
+	return out
+}
 
-	// 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
-		}
-		// 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})
+func testRobustness() {
+	ctx := context.Background()
+	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()
+
+			etcdTraffic.RunTrafficLoop(ctx, c, limiter,
+				identity.NewIDProvider(),
+				identity.NewLeaseIDStorage(),
+				traffic.NewConcurrencyLimiter(profile.MaxNonUniqueRequestConcurrency),
+				finish,
+			)
+			mux.Lock()
+			reports = append(reports, c.Report())
+			mux.Unlock()
+		}(c)
 	}
-	cli.Close()
-
-	assert.Reachable("Completion of a key deleting check", nil)
-	log.Printf("Completion of a key deleting check")
 }
 
 func main() {
-	DeleteKeys()
+	testRobustness()
 }

tests/robustness/traffic/etcd.go

@@ -31,12 +31,12 @@ import (
 )
 
 var (
-	EtcdPutDeleteLease Traffic = etcdTraffic{
-		keyCount:     10,
-		leaseTTL:     DefaultLeaseTTL,
-		largePutSize: 32769,
+	EtcdPutDeleteLease Traffic = EtcdTraffic{
+		KeyCount:     10,
+		LeaseTTL:     DefaultLeaseTTL,
+		LargePutSize: 32769,
 		// Please keep the sum of weights equal 100.
-		requests: []random.ChoiceWeight[etcdRequestType]{
+		Requests: []random.ChoiceWeight[EtcdRequestType]{
 			{Choice: Get, Weight: 15},
 			{Choice: List, Weight: 15},
 			{Choice: StaleGet, Weight: 10},
@@ -50,12 +50,12 @@ var (
 			{Choice: LargePut, Weight: 5},
 		},
 	}
-	EtcdPut Traffic = etcdTraffic{
-		keyCount:     10,
-		largePutSize: 32769,
-		leaseTTL:     DefaultLeaseTTL,
+	EtcdPut Traffic = EtcdTraffic{
+		KeyCount:     10,
+		LargePutSize: 32769,
+		LeaseTTL:     DefaultLeaseTTL,
 		// Please keep the sum of weights equal 100.
-		requests: []random.ChoiceWeight[etcdRequestType]{
+		Requests: []random.ChoiceWeight[EtcdRequestType]{
 			{Choice: Get, Weight: 15},
 			{Choice: List, Weight: 15},
 			{Choice: StaleGet, Weight: 10},
@@ -65,56 +65,56 @@ var (
 			{Choice: Put, Weight: 40},
 		},
 	}
-	EtcdDelete Traffic = etcdTraffic{
-		keyCount:     10,
-		largePutSize: 32769,
-		leaseTTL:     DefaultLeaseTTL,
+	EtcdDelete Traffic = EtcdTraffic{
+		KeyCount:     10,
+		LargePutSize: 32769,
+		LeaseTTL:     DefaultLeaseTTL,
 		// Please keep the sum of weights equal 100.
-		requests: []random.ChoiceWeight[etcdRequestType]{
+		Requests: []random.ChoiceWeight[EtcdRequestType]{
 			{Choice: Put, Weight: 50},
 			{Choice: Delete, Weight: 50},
 		},
 	}
 )
 
-type etcdTraffic struct {
-	keyCount     int
-	requests     []random.ChoiceWeight[etcdRequestType]
-	leaseTTL     int64
-	largePutSize int
+type EtcdTraffic struct {
+	KeyCount     int
+	Requests     []random.ChoiceWeight[EtcdRequestType]
+	LeaseTTL     int64
+	LargePutSize int
 }
 
-func (t etcdTraffic) ExpectUniqueRevision() bool {
+func (t EtcdTraffic) ExpectUniqueRevision() bool {
 	return false
 }
 
-type etcdRequestType string
+type EtcdRequestType string
 
 const (
-	Get           etcdRequestType = "get"
-	StaleGet      etcdRequestType = "staleGet"
-	List          etcdRequestType = "list"
-	StaleList     etcdRequestType = "staleList"
-	Put           etcdRequestType = "put"
-	LargePut      etcdRequestType = "largePut"
-	Delete        etcdRequestType = "delete"
-	MultiOpTxn    etcdRequestType = "multiOpTxn"
-	PutWithLease  etcdRequestType = "putWithLease"
-	LeaseRevoke   etcdRequestType = "leaseRevoke"
-	CompareAndSet etcdRequestType = "compareAndSet"
-	Defragment    etcdRequestType = "defragment"
+	Get           EtcdRequestType = "get"
+	StaleGet      EtcdRequestType = "staleGet"
+	List          EtcdRequestType = "list"
+	StaleList     EtcdRequestType = "staleList"
+	Put           EtcdRequestType = "put"
+	LargePut      EtcdRequestType = "largePut"
+	Delete        EtcdRequestType = "delete"
+	MultiOpTxn    EtcdRequestType = "multiOpTxn"
+	PutWithLease  EtcdRequestType = "putWithLease"
+	LeaseRevoke   EtcdRequestType = "leaseRevoke"
+	CompareAndSet EtcdRequestType = "compareAndSet"
+	Defragment    EtcdRequestType = "defragment"
 )
 
-func (t etcdTraffic) Name() string {
+func (t EtcdTraffic) Name() string {
 	return "Etcd"
 }
 
-func (t etcdTraffic) RunTrafficLoop(ctx context.Context, c *client.RecordingClient, limiter *rate.Limiter, ids identity.Provider, lm identity.LeaseIDStorage, nonUniqueWriteLimiter ConcurrencyLimiter, finish <-chan struct{}) {
+func (t EtcdTraffic) RunTrafficLoop(ctx context.Context, c *client.RecordingClient, limiter *rate.Limiter, ids identity.Provider, lm identity.LeaseIDStorage, nonUniqueWriteLimiter ConcurrencyLimiter, finish <-chan struct{}) {
 	lastOperationSucceeded := true
 	var lastRev int64
-	var requestType etcdRequestType
+	var requestType EtcdRequestType
 	client := etcdTrafficClient{
-		etcdTraffic:  t,
+		EtcdTraffic:  t,
 		keyPrefix:    "key",
 		client:       c,
 		limiter:      limiter,
@@ -133,7 +133,7 @@ func (t etcdTraffic) RunTrafficLoop(ctx context.Context, c *client.RecordingClie
 
 		// Avoid multiple failed writes in a row
 		if lastOperationSucceeded {
-			choices := t.requests
+			choices := t.Requests
 			if shouldReturn = nonUniqueWriteLimiter.Take(); !shouldReturn {
 				choices = filterOutNonUniqueEtcdWrites(choices)
 			}
@@ -156,7 +156,7 @@ func (t etcdTraffic) RunTrafficLoop(ctx context.Context, c *client.RecordingClie
 	}
 }
 
-func (t etcdTraffic) RunCompactLoop(ctx context.Context, c *client.RecordingClient, period time.Duration, finish <-chan struct{}) {
+func (t EtcdTraffic) RunCompactLoop(ctx context.Context, c *client.RecordingClient, period time.Duration, finish <-chan struct{}) {
 	var lastRev int64 = 2
 	ticker := time.NewTicker(period)
 	defer ticker.Stop()
@@ -185,7 +185,7 @@ func (t etcdTraffic) RunCompactLoop(ctx context.Context, c *client.RecordingClie
 	}
 }
 
-func filterOutNonUniqueEtcdWrites(choices []random.ChoiceWeight[etcdRequestType]) (resp []random.ChoiceWeight[etcdRequestType]) {
+func filterOutNonUniqueEtcdWrites(choices []random.ChoiceWeight[EtcdRequestType]) (resp []random.ChoiceWeight[EtcdRequestType]) {
 	for _, choice := range choices {
 		if choice.Choice != Delete && choice.Choice != LeaseRevoke {
 			resp = append(resp, choice)
@@ -195,15 +195,15 @@ func filterOutNonUniqueEtcdWrites(choices []random.ChoiceWeight[etcdRequestType]
 }
 
 type etcdTrafficClient struct {
-	etcdTraffic
+	EtcdTraffic
 	keyPrefix    string
 	client       *client.RecordingClient
 	limiter      *rate.Limiter
 	idProvider   identity.Provider
 	leaseStorage identity.LeaseIDStorage
 }
 
-func (c etcdTrafficClient) Request(ctx context.Context, request etcdRequestType, lastRev int64) (rev int64, err error) {
+func (c etcdTrafficClient) Request(ctx context.Context, request EtcdRequestType, lastRev int64) (rev int64, err error) {
 	opCtx, cancel := context.WithTimeout(ctx, RequestTimeout)
 	defer cancel()
 
@@ -241,7 +241,7 @@ func (c etcdTrafficClient) Request(ctx context.Context, request etcdRequestType,
 		}
 	case LargePut:
 		var resp *clientv3.PutResponse
-		resp, err = c.client.Put(opCtx, c.randomKey(), random.RandString(c.largePutSize))
+		resp, err = c.client.Put(opCtx, c.randomKey(), random.RandString(c.LargePutSize))
 		if resp != nil {
 			rev = resp.Header.Revision
 		}
@@ -290,7 +290,7 @@ func (c etcdTrafficClient) Request(ctx context.Context, request etcdRequestType,
 		leaseID := c.leaseStorage.LeaseID(c.client.ID)
 		if leaseID == 0 {
 			var resp *clientv3.LeaseGrantResponse
-			resp, err = c.client.LeaseGrant(opCtx, c.leaseTTL)
+			resp, err = c.client.LeaseGrant(opCtx, c.LeaseTTL)
 			if resp != nil {
 				leaseID = int64(resp.ID)
 				rev = resp.ResponseHeader.Revision
@@ -335,7 +335,7 @@ func (c etcdTrafficClient) Request(ctx context.Context, request etcdRequestType,
 }
 
 func (c etcdTrafficClient) pickMultiTxnOps() (ops []clientv3.Op) {
-	keys := rand.Perm(c.keyCount)
+	keys := rand.Perm(c.KeyCount)
 	opTypes := make([]model.OperationType, 4)
 
 	atLeastOnePut := false
@@ -372,10 +372,10 @@ func (c etcdTrafficClient) randomKey() string {
 }
 
 func (c etcdTrafficClient) key(i int) string {
-	return fmt.Sprintf("%s%d", c.keyPrefix, i%c.keyCount)
+	return fmt.Sprintf("%s%d", c.keyPrefix, i%c.KeyCount)
 }
 
-func (t etcdTraffic) pickOperationType() model.OperationType {
+func (t EtcdTraffic) pickOperationType() model.OperationType {
 	roll := rand.Int() % 100
 	if roll < 10 {
 		return model.DeleteOperation