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

Code before:
		return
	}
	if result.Linearization.Linearizable == porcupine.Unknown {
		assert.Unreachable("Linearization timeout", nil)
	} else {
		assert.Always(result.Linearization.Linearizable == porcupine.Ok, "Linearization validation passes", nil)
	}
	r.Visualize = result.Linearization.Visualize
	assert.Always(result.WatchError == nil, "Watch validation passes", map[string]any{"error": result.WatchError})
	assert.Always(result.SerializableError == nil, "Serializable validation passes", map[string]any{"error": result.WatchError})
	lg.Info("Completed robustness validation")
}

func runTraffic(ctx context.Context, lg *zap.Logger, hosts []string, baseTime time.Time, duration time.Duration) []report.ClientReport {
	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)
	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

Review comment (suggestion):
Let's keep this function simple. This loop should not be needed if we only passed hosts to `CollectClusterWatchEvents`. This function doesn't need to accept `etcdEndpoints [][]string`, just `hosts []string` would be enough

Code after:
	validateConfig := validate.Config{ExpectRevisionUnique: traffic.EtcdAntithesis.ExpectUniqueRevision()}
	result, err := validate.ValidateAndReturnVisualize(lg, validateConfig, r.Client, etcdDataDirs, 5*time.Minute)
	if err != nil {
		lg.Info("Validation error", zap.Error(err))
		assert.Unreachable("Validation error", map[string]any{"error": err})
		return
	}
	if result.Linearization.Linearizable == porcupine.Unknown {
		assert.Unreachable("Linearization timeout", nil)
	} else {
		assert.Always(result.Linearization.Linearizable == porcupine.Ok, "Linearization validation passes", nil)
	}
	r.Visualize = result.Linearization.Visualize
	assert.Always(result.WatchError == nil, "Watch validation passes", map[string]any{"error": result.WatchError})
	assert.Always(result.SerializableError == nil, "Serializable validation passes", map[string]any{"error": result.WatchError})
	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
}