Repository: etcd-io/etcd
File: tests/antithesis/test-template/robustness/main.go (Go)

Code before:
	r, err := traffic.CheckEmptyDatabaseAtStart(ctx, lg, hosts, ids, baseTime)
	if err != nil {
		lg.Fatal("Failed empty database at start check", zap.Error(err))
	}
	reports := []report.ClientReport{r}
	watchReport := []report.ClientReport{}
	maxRevisionChan := make(chan int64, 1)
	endpoints := make([][]string, len(hosts))
	for _, host := range hosts {
		endpoints = append(endpoints, []string{host})
	}
	watchConfig := client.WatchConfig{
		RequestProgress: true,
	}
	g := errgroup.Group{}
	g.Go(func() error {
		defer close(maxRevisionChan)
		reports := slices.Concat(reports, simulateTraffic(ctx, hosts, baseTime, duration))
		maxRevision := report.OperationsMaxRevision(reports)
		maxRevisionChan <- maxRevision
		lg.Info("Finished simulating Traffic", zap.Int64("max-revision", maxRevision))
		return nil
	})
	g.Go(func() error {
		var watchErr error
		watchReport, _, watchErr = client.CollectClusterWatchEvents(ctx, lg, endpoints, maxRevisionChan, watchConfig, time.Now(), ids)
		return watchErr
	})
	g.Wait()

	return slices.Concat(reports, watchReport)
}

func simulateTraffic(ctx context.Context, hosts []string, baseTime time.Time, duration time.Duration) []report.ClientReport {
	var mux sync.Mutex
	var wg sync.WaitGroup
	ids := identity.NewIDProvider()
	storage := identity.NewLeaseIDStorage()
	limiter := rate.NewLimiter(rate.Limit(profile.MaximalQPS), profile.BurstableQPS)
	concurrencyLimiter := traffic.NewConcurrencyLimiter(profile.MaxNonUniqueRequestConcurrency)
	finish := closeAfter(ctx, duration)
	reports := []report.ClientReport{}
	for i := 0; i < profile.ClientCount; i++ {
		c := connect([]string{hosts[i%len(hosts)]}, ids, baseTime)
		defer c.Close()
		wg.Add(1)
		go func(c *client.RecordingClient) {
			defer wg.Done()
			defer c.Close()

			traffic.EtcdAntithesis.RunTrafficLoop(ctx, c, limiter,

Review comment (suggestion):
```suggestion
		watchReport, _, watchErr = client.CollectClusterWatchEvents(ctx, lg, endpoints, maxRevisionChan, watchConfig, baseTime, ids)
```

Code after:
	lg.Info("Completed robustness validation")
}

func runTraffic(ctx context.Context, lg *zap.Logger, hosts []string, baseTime time.Time, duration time.Duration) ([]report.ClientReport, error) {
	ids := identity.NewIDProvider()
	r, err := traffic.CheckEmptyDatabaseAtStart(ctx, lg, hosts, ids, baseTime)
	if err != nil {
		lg.Fatal("Failed empty database at start check", zap.Error(err))
	}
	reports := []report.ClientReport{r}
	watchReport := []report.ClientReport{}
	maxRevisionChan := make(chan int64, 1)
	watchConfig := client.WatchConfig{
		RequestProgress: true,
	}
	g := errgroup.Group{}
	g.Go(func() error {
		defer close(maxRevisionChan)
		reports := slices.Concat(reports, simulateTraffic(ctx, hosts, ids, baseTime, duration))
		maxRevision := report.OperationsMaxRevision(reports)
		maxRevisionChan <- maxRevision
		lg.Info("Finished simulating Traffic", zap.Int64("max-revision", maxRevision))
		return nil
	})
	g.Go(func() error {
		var watchErr error
		watchReport, _, watchErr = client.CollectClusterWatchEvents(ctx, lg, hosts, maxRevisionChan, watchConfig, baseTime, ids)
		return watchErr
	})
	if err := g.Wait(); err != nil {
		return nil, err
	}

	return slices.Concat(reports, watchReport), nil
}

func simulateTraffic(ctx context.Context, hosts []string, ids identity.Provider, baseTime time.Time, duration time.Duration) []report.ClientReport {
	var mux sync.Mutex
	var wg sync.WaitGroup
	storage := identity.NewLeaseIDStorage()
	limiter := rate.NewLimiter(rate.Limit(profile.MaximalQPS), profile.BurstableQPS)
	concurrencyLimiter := traffic.NewConcurrencyLimiter(profile.MaxNonUniqueRequestConcurrency)
	finish := closeAfter(ctx, duration)
	reports := []report.ClientReport{}
	for i := 0; i < profile.ClientCount; i++ {
		c := connect([]string{hosts[i%len(hosts)]}, ids, baseTime)
		defer c.Close()
		wg.Add(1)
		go func(c *client.RecordingClient) {
			defer wg.Done()
			defer c.Close()