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

Code before:
	lg.Info("Start traffic generation", zap.Duration("duration", duration))
	r := report.TestReport{Logger: lg}

	lg.Info("Start traffic generation")
	r.Client = runTraffic(ctx, lg, hosts, baseTime, duration)
	lg.Info("Completed traffic generation")

	etcdDataDirs, err := report.PersistedRequestsDirs(lg, persistedRequestdirs)
	if err != nil {
		lg.Error("Failed to create PersistedRequestdirs to access etcd WAL files")
		panic(err)
	}
	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
	if err := r.Report(true, false, true, reportPath); err != nil {
		lg.Error("Failed to save validation report", zap.Error(err))
	}
	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 {
	limiter := rate.NewLimiter(rate.Limit(profile.MaximalQPS), profile.BurstableQPS)
	ids := identity.NewIDProvider()
	storage := identity.NewLeaseIDStorage()
	concurrencyLimiter := traffic.NewConcurrencyLimiter(profile.MaxNonUniqueRequestConcurrency)

	r, err := traffic.CheckEmptyDatabaseAtStart(ctx, lg, hosts, ids, baseTime)
	if err != nil {
		lg.Fatal("Failed empty database at start check", zap.Error(err))
	}

	finish := closeAfter(ctx, duration)
	reports := []report.ClientReport{r}
	var mux sync.Mutex
	var wg sync.WaitGroup
	for i := 0; i < profile.ClientCount; i++ {
		c := connect([]string{hosts[i%len(hosts)]}, ids, baseTime)
		defer c.Close()

Review comment (suggestion):
```suggestion
	if err := r.Report(reportPath); err != nil {
```

Code after:
	r := report.TestReport{Logger: lg, ServersDataPath: map[string]string{}}
	for i := range len(persistedRequestdirs) {
		r.ServersDataPath[fmt.Sprintf("etcd%d", i)] = persistedRequestdirs[i]
	}
	defer func() {
		if err := r.Report(reportPath); err != nil {
			lg.Error("Failed to save validation report", zap.Error(err))
		}
	}()
	testRobustness(ctx, lg, &r, hosts, baseTime, duration)
}

func testRobustness(ctx context.Context, lg *zap.Logger, r *report.TestReport, hosts []string, baseTime time.Time, duration time.Duration) {
	lg.Info("Start traffic generation", zap.Duration("duration", duration))
	r.Client = runTraffic(ctx, lg, hosts, baseTime, duration)
	lg.Info("Completed traffic generation")

	etcdDataDirs, err := report.PersistedRequestsDirs(lg, slices.Collect(maps.Values(r.ServersDataPath)))
	if err != nil {
		lg.Error("Failed to create PersistedRequestdirs to access etcd WAL files")
		panic(err)
	}
	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 {
	limiter := rate.NewLimiter(rate.Limit(profile.MaximalQPS), profile.BurstableQPS)
	ids := identity.NewIDProvider()
	storage := identity.NewLeaseIDStorage()
	concurrencyLimiter := traffic.NewConcurrencyLimiter(profile.MaxNonUniqueRequestConcurrency)

	r, err := traffic.CheckEmptyDatabaseAtStart(ctx, lg, hosts, ids, baseTime)
	if err != nil {
		lg.Fatal("Failed empty database at start check", zap.Error(err))
	}