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

Code before:
	})
	watchSet := client.NewSet(ids, baseTime)
	defer watchSet.Close()
	g.Go(func() error {
		err := client.CollectClusterWatchEvents(ctx, lg, hosts, maxRevisionChan, watchConfig, watchSet)
		return err
	})
	if err := g.Wait(); err != nil {
		return nil, err
	}
	endTime := time.Since(baseTime)
	reports := slices.Concat(trafficSet.Reports(), watchSet.Reports())
	totalStats := traffic.CalculateStats(reports, startTime, endTime)
	lg.Info("Completed traffic generation",
		zap.Int("successes", totalStats.Successes),
		zap.Int("failures", totalStats.Failures),
		zap.Float64("successRate", totalStats.SuccessRate()),
		zap.Duration("period", totalStats.Period),
		zap.Float64("qps", totalStats.QPS()),
	)
	return reports, nil
}

func simulateTraffic(ctx context.Context, tf traffic.Traffic, hosts []string, clientSet *client.ClientSet, duration time.Duration) {
	var wg sync.WaitGroup
	defer clientSet.Close()
	storage := identity.NewLeaseIDStorage()
	limiter := rate.NewLimiter(rate.Limit(profile.MaximalQPS), profile.BurstableQPS)
	concurrencyLimiter := traffic.NewConcurrencyLimiter(profile.MaxNonUniqueRequestConcurrency)
	finish := closeAfter(ctx, duration)
	keyStore := traffic.NewKeyStore(10, "key")
	for i := range profile.ClientCount {
		c := connect(clientSet, []string{hosts[i%len(hosts)]})
		wg.Add(1)
		defer c.Close()
		go func(c *client.RecordingClient) {
			defer wg.Done()
			tf.RunTrafficLoop(ctx, c, limiter,
				clientSet.IdentityProvider(),
				storage,
				concurrencyLimiter,
				keyStore,
				finish,
			)
		}(c)
	}
	wg.Add(1)
	compactClient := connect(clientSet, hosts)
	go func(c *client.RecordingClient) {
		defer wg.Done()
		tf.RunCompactLoop(ctx, c, traffic.DefaultCompactionPeriod, finish)

Review comment (suggestion):
simulateTraffic function didn't create the clientSet thus it should not be responsible for closing it.

Code after:
	})
	watchSet := client.NewSet(ids, baseTime)
	defer watchSet.Close()
	g.Go(func() error {
		err := client.CollectClusterWatchEvents(ctx, lg, hosts, maxRevisionChan, watchConfig, watchSet)
		return err
	})
	if err := g.Wait(); err != nil {
		return nil, err
	}
	endTime := time.Since(baseTime)
	reports := slices.Concat(trafficSet.Reports(), watchSet.Reports())
	totalStats := traffic.CalculateStats(reports, startTime, endTime)
	lg.Info("Completed traffic generation",
		zap.Int("successes", totalStats.Successes),
		zap.Int("failures", totalStats.Failures),
		zap.Float64("successRate", totalStats.SuccessRate()),
		zap.Duration("period", totalStats.Period),
		zap.Float64("qps", totalStats.QPS()),
	)
	return reports, nil
}

func simulateTraffic(ctx context.Context, tf traffic.Traffic, hosts []string, clientSet *client.ClientSet, duration time.Duration) {
	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)
	keyStore := traffic.NewKeyStore(10, "key")
	for i := range profile.ClientCount {
		c := connect(clientSet, []string{hosts[i%len(hosts)]})
		wg.Add(1)
		go func(c *client.RecordingClient) {
			defer wg.Done()
			defer c.Close()
			tf.RunTrafficLoop(ctx, c, limiter,
				clientSet.IdentityProvider(),
				storage,
				concurrencyLimiter,
				keyStore,
				finish,
			)
		}(c)
	}
	wg.Add(1)
	compactClient := connect(clientSet, hosts)
	go func(c *client.RecordingClient) {
		defer wg.Done()
		defer c.Close()
		tf.RunCompactLoop(ctx, c, traffic.DefaultCompactionPeriod, finish)