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

Code before:
				defer c.Close()

				traffic.EtcdAntithesis.RunTrafficLoop(ctx, c, limiter,
					ids,
					storage,
					concurrencyLimiter,
					finish,
				)
				mux.Lock()
				reports = append(reports, c.Report())
				mux.Unlock()
			}(c)
		}
		wg.Wait()
		maxRevision := report.OperationsMaxRevision(reports)
		maxRevisionChan <- maxRevision
		lg.Info("Finished simulating Traffic", zap.Int64("max-revision", maxRevision))
	}()
	go func() {
		endpoints := make([][]string, len(hosts))
		for _, host := range hosts {
			endpoints = append(endpoints, []string{host})
		}
		watchConfig := client.WatchConfig{
			RequestProgress: true,
		}
		var watchErr error
		watchReport, _, watchErr = client.CollectClusterWatchEvents(ctx, lg, endpoints, maxRevisionChan, watchConfig, time.Now(), ids)
		if watchErr != nil {
			panic(watchErr)
		}
	}()
	return append(reports, watchReport...)
}

func connect(endpoints []string, ids identity.Provider, baseTime time.Time) *client.RecordingClient {
	cli, err := client.NewRecordingClient(endpoints, 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{"endpoints": endpoints, "error": err})
		os.Exit(1)
	}
	return cli
}

func closeAfter(ctx context.Context, t time.Duration) <-chan struct{} {
	out := make(chan struct{})
	go func() {
		select {
		case <-time.After(t):
		case <-ctx.Done():

Review comment (refactor):
Move this code outside goroutine. Goroutines should have minimal code needed for coordination between them.

Code after:
		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()

			traffic.EtcdAntithesis.RunTrafficLoop(ctx, c, limiter,
				ids,
				storage,
				concurrencyLimiter,
				finish,
			)
			mux.Lock()
			reports = append(reports, c.Report())
			mux.Unlock()
		}(c)
	}
	wg.Wait()
	return reports
}

func connect(endpoints []string, ids identity.Provider, baseTime time.Time) *client.RecordingClient {
	cli, err := client.NewRecordingClient(endpoints, 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{"endpoints": endpoints, "error": err})
		os.Exit(1)