Merge #153584

153584: asim: trace mma.ComputeChanges and save to generated artifacts r=tbg a=tbg

This makes it much easier to figure out what a particular store's MMA was doing.
Traces will be in `generated/<testname>/traces/<store>`.

Epic: CRDB-49117

Co-authored-by: Tobias Grieger <tobias.b.grieger@gmail.com>

Author: craig[bot]

pkg/kv/kvserver/asim/BUILD.bazel

@@ -18,6 +18,8 @@ go_library(
         "//pkg/kv/kvserver/asim/storerebalancer",
         "//pkg/kv/kvserver/asim/workload",
         "//pkg/util/log",
+        "//pkg/util/tracing",
+        "//pkg/util/tracing/tracingpb",
     ],
 )
 

pkg/kv/kvserver/asim/asim.go

@@ -21,12 +21,16 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/storerebalancer"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/workload"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
 )
 
 // Simulator simulates an entire cluster, and runs the allocator of each store
 // in that cluster.
 type Simulator struct {
 	log.AmbientContext
+	onRecording func(storeID state.StoreID, rec tracingpb.Recording)
+
 	curr time.Time
 	end  time.Time
 	// interval is the step between ticks for active simulaton components, such
@@ -96,22 +100,27 @@ func NewSimulator(
 
 	s := &Simulator{
 		AmbientContext: log.MakeTestingAmbientCtxWithNewTracer(),
-		curr:           settings.StartTime,
-		end:            settings.StartTime.Add(duration),
-		interval:       settings.TickInterval,
-		generators:     wgs,
-		state:          initialState,
-		changer:        changer,
-		rqs:            rqs,
-		lqs:            lqs,
-		sqs:            sqs,
-		controllers:    controllers,
-		srs:            srs,
-		mmSRs:          mmSRs,
-		pacers:         pacers,
-		gossip:         gossip.NewGossip(initialState, settings),
-		metrics:        m,
-		shuffler:       state.NewShuffler(settings.Seed),
+		onRecording: func(storeID state.StoreID, rec tracingpb.Recording) {
+			if fn := settings.OnRecording; fn != nil {
+				fn(int64(storeID), rec)
+			}
+		},
+		curr:        settings.StartTime,
+		end:         settings.StartTime.Add(duration),
+		interval:    settings.TickInterval,
+		generators:  wgs,
+		state:       initialState,
+		changer:     changer,
+		rqs:         rqs,
+		lqs:         lqs,
+		sqs:         sqs,
+		controllers: controllers,
+		srs:         srs,
+		mmSRs:       mmSRs,
+		pacers:      pacers,
+		gossip:      gossip.NewGossip(initialState, settings),
+		metrics:     m,
+		shuffler:    state.NewShuffler(settings.Seed),
 		// TODO(kvoli): Keeping the state around is a bit hacky, find a better
 		// method of reporting the ranges.
 		history:       history.History{Recorded: [][]metrics.StoreMetrics{}, S: initialState},
@@ -380,7 +389,14 @@ func (s *Simulator) tickMMStoreRebalancers(ctx context.Context, tick time.Time,
 	stores := s.state.Stores()
 	s.shuffler(len(stores), func(i, j int) { stores[i], stores[j] = stores[j], stores[i] })
 	for _, store := range stores {
+		var finishAndGetRecording func() tracingpb.Recording
+		if s.onRecording != nil {
+			ctx, finishAndGetRecording = tracing.ContextWithRecordingSpan(ctx, s.Tracer, "mma.ComputeChanges")
+		}
 		s.mmSRs[store.StoreID()].Tick(ctx, tick, state)
+		if finishAndGetRecording != nil {
+			s.onRecording(store.StoreID(), finishAndGetRecording())
+		}
 	}
 }
 

pkg/kv/kvserver/asim/config/BUILD.bazel

@@ -5,5 +5,8 @@ go_library(
     srcs = ["settings.go"],
     importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config",
     visibility = ["//visibility:public"],
-    deps = ["//pkg/settings/cluster"],
+    deps = [
+        "//pkg/settings/cluster",
+        "//pkg/util/tracing/tracingpb",
+    ],
 )

pkg/kv/kvserver/asim/config/settings.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
 )
 
 const (
@@ -116,6 +117,9 @@ type SimulationSettings struct {
 	// TODO(wenyihu6): Remove any non-simulation settings from this struct and
 	// instead override the settings below.
 	ST *cluster.Settings
+	// OnRecording is called with trace spans obtained by recording the allocator.
+	// NB: we can't use state.StoreID here since that causes an import cycle.
+	OnRecording func(storeID int64, rec tracingpb.Recording)
 }
 
 // DefaultSimulationSettings returns a set of default settings for simulation.

pkg/kv/kvserver/asim/gen/printer.go

@@ -33,7 +33,8 @@ func compareSettingsToDefault(settings config.SimulationSettings) string {
 		// Skip cluster setting ST and seed. The simulation seed is derived from
 		// rand.New(rand.NewSource(42)).Int63() by default, while the default
 		// simulation setting itself uses 42.
-		if field.Name == "ST" || field.Name == "Seed" {
+		// We also skip the OnRecording callback.
+		if field.Name == "ST" || field.Name == "Seed" || field.Name == "OnRecording" {
 			continue
 		}
 

pkg/kv/kvserver/asim/tests/BUILD.bazel

@@ -58,8 +58,10 @@ go_test(
         "//pkg/util/humanizeutil",
         "//pkg/util/leaktest",
         "//pkg/util/log",
+        "//pkg/util/tracing/tracingpb",
         "@com_github_cockroachdb_datadriven//:datadriven",
         "@com_github_cockroachdb_logtags//:logtags",
+        "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/kv/kvserver/asim/tests/datadriven_simulation_test.go

@@ -8,11 +8,11 @@ package tests
 import (
 	"context"
 	"fmt"
-	"hash"
 	"hash/fnv"
 	"math/rand"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"testing"
 	"time"
@@ -35,8 +35,10 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
 	"github.com/cockroachdb/datadriven"
 	"github.com/cockroachdb/logtags"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -165,6 +167,12 @@ var runAsimTests = envutil.EnvOrDefaultBool("COCKROACH_RUN_ASIM_TESTS", false)
 //     random number generator that creates the seed used to generate each
 //     simulation sample. The default values are: duration=30m (30 minutes)
 //     samples=1 seed=random.
+//
+// To run all tests and rewrite the testdata files as well as generate the
+// artifacts in `testdata/generated`, you can use:
+/*
+./dev test pkg/kv/kvserver/asim/tests --ignore-cache --rewrite -v -f TestDataDriven -- --test_env COCKROACH_RUN_ASIM_TESTS=true --test_env COCKROACH_ALWAYS_KEEP_TEST_LOGS=true
+*/
 func TestDataDriven(t *testing.T) {
 	skip.UnderDuressWithIssue(t, 149875)
 	leakTestAfter := leaktest.AfterTest(t)
@@ -531,7 +539,34 @@ func TestDataDriven(t *testing.T) {
 						require.NotNil(t, set, "unknown mode value: %s", mv)
 						set(&eventGen)
 
+						// TODO(tbg): need to decide whether multiple evals in a single file
+						// is a feature or an anti-pattern. If it's a feature, we should let
+						// the `name` part below be adjustable (but not the plotDir) via a
+						// parameter to the `eval` command.
+						testName := name + "_" + mv
+
 						for sample := 0; sample < samples; sample++ {
+							recIdx := map[int64]int{}
+							settingsGen.Settings.OnRecording = func(storeID int64, rec tracingpb.Recording) {
+								if !rewrite || len(rec[0].Logs) == 0 {
+									return
+								}
+								traceDir := filepath.Join(plotDir, "traces", fmt.Sprintf("s%d", storeID))
+								if recIdx[storeID] == 0 {
+									require.NoError(t, os.MkdirAll(traceDir, 0755))
+								}
+								re := regexp.MustCompile(`[^a-zA-Z0-9]+`)
+								outName := fmt.Sprintf("%s_%s_s%d", mv, re.ReplaceAllString(rec[0].Operation, "_"), storeID)
+								if sample > 0 {
+									outName += fmt.Sprintf("_sample%d", sample+1)
+								}
+								outName += "_" + fmt.Sprintf("%03d.txt", recIdx[storeID])
+								assert.NoError(t, os.WriteFile(
+									filepath.Join(traceDir, outName),
+									[]byte(rec.String()), 0644))
+								recIdx[storeID] += 1
+							}
+
 							assertionFailures := []string{}
 							var tmpStrB *strings.Builder = nil
 							if stateStrForOnce == "" {
@@ -561,17 +596,10 @@ func TestDataDriven(t *testing.T) {
 						// Generate artifacts. Hash artifact input data to ensure they are
 						// up to date.
 						hasher := fnv.New64a()
-						// TODO(tbg): need to decide whether multiple evals in a single file
-						// is a feature or an anti-pattern. If it's a feature, we should let
-						// the `name` part below be adjustable (but not the plotDir) via a
-						// parameter to the `eval` command.
-						testName := name + "_" + mv
+
 						for sample, h := range run.hs {
-							generateAllPlots(t, &buf, h, testName, sample+1, plotDir, hasher, rewrite,
+							printStatsAndGenerateJSON(t, &buf, h, testName, sample+1, plotDir, hasher, rewrite,
 								settingsGen.Settings.TickInterval, metricsMap)
-							generateTopology(t, h,
-								filepath.Join(plotDir, fmt.Sprintf("%s_%d_topology.txt", testName, sample+1)),
-								hasher, rewrite)
 						}
 						artifactsHash := hasher.Sum64()
 
@@ -707,22 +735,6 @@ type modeHistory struct {
 	hs   []history.History
 }
 
-func generateTopology(
-	t *testing.T, h history.History, topFile string, hasher hash.Hash, rewrite bool,
-) {
-	// TODO(tbg): this can in principle be printed without even
-	// evaluating the test, and in particular it's independent of
-	// settings. It seems like an artifact of the implementation
-	// that we can only access the structured topology after the
-	// simulation has run.
-	top := h.S.Topology()
-	s := top.String()
-	_, _ = fmt.Fprint(hasher, s)
-	if rewrite {
-		require.NoError(t, os.WriteFile(topFile, []byte(s), 0644))
-	}
-}
-
 // writeStateStrToFile writes the state string to the given file.
 func writeStateStrToFile(t *testing.T, topFile string, stateStr string, rewrite bool) {
 	if rewrite {

pkg/kv/kvserver/asim/tests/plotting_test.go

@@ -21,14 +21,13 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// generateAllPlots creates plots for all available metrics from simulation history.
+// printStatsAndGenerateJSON creates plots for all available metrics from simulation history.
 // All plot files are hashed directly to the provided hasher.
 // If rewrite is false, plots are generated but not saved to disk.
 // Returns a slice of filenames for all generated plots.
 //
 // TODO(tbg): introduce a SimulationEnv and make this a method on it.
-// TODO(tbg): rename to generateArtifacts.
-func generateAllPlots(
+func printStatsAndGenerateJSON(
 	t *testing.T,
 	buf *strings.Builder,
 	h history.History,

pkg/kv/kvserver/asim/tests/testdata/generated/.gitignore

@@ -4,6 +4,9 @@
 !*/
 
 # Except this specific directory (relative to .gitignore)
+# Note that this intentionally does not include subdirectories (such as traces)
+# which are not deterministic.
 !example_rebalancing/**
+**/traces
 # And the readme file.
-!README.md
+!README.md

pkg/kv/kvserver/asim/tests/testdata/non_rand/decommission_conformance.txt

@@ -28,7 +28,7 @@ assertion type=conformance under=0 over=0 unavailable=0 violating=0
 
 eval duration=20m cfgs=(sma-count,mma-only)
 ----
-artifacts[sma-count]: 19128202c9e746ab
+artifacts[sma-count]: 8e9d6ada04cb9561
 ==========================
-artifacts[mma-only]: d80019be362b72f7
+artifacts[mma-only]: 239144d63ffded95
 ==========================