asim: further improve tracing

Declutter the main datadriven test by extracting out a tracing helper. Revamp
the file format so that it includes the seconds elapsed since simulation start,
which facilitates lining traces up with asimview.

We now also trace lease and replica count rebalancing decisions, as well as any
constraint repair.

Epic: CRDB-49117

Author: tbg

pkg/kv/kvserver/asim/asim.go

@@ -29,7 +29,7 @@ import (
 // in that cluster.
 type Simulator struct {
 	log.AmbientContext
-	onRecording func(storeID state.StoreID, rec tracingpb.Recording)
+	onRecording func(storeID state.StoreID, atDuration time.Duration, rec tracingpb.Recording)
 
 	curr time.Time
 	end  time.Time
@@ -100,9 +100,9 @@ func NewSimulator(
 
 	s := &Simulator{
 		AmbientContext: log.MakeTestingAmbientCtxWithNewTracer(),
-		onRecording: func(storeID state.StoreID, rec tracingpb.Recording) {
+		onRecording: func(storeID state.StoreID, atDuration time.Duration, rec tracingpb.Recording) {
 			if fn := settings.OnRecording; fn != nil {
-				fn(int64(storeID), rec)
+				fn(int64(storeID), atDuration, rec)
 			}
 		},
 		curr:        settings.StartTime,
@@ -323,6 +323,25 @@ func (s *Simulator) tickStoreClocks(tick time.Time) {
 	s.state.TickClock(tick)
 }
 
+func (s *Simulator) doAndMaybeTrace(
+	ctx context.Context,
+	storeID state.StoreID,
+	tick time.Time,
+	op string,
+	f func(ctx context.Context),
+) {
+	atDuration := tick.Sub(s.settings.StartTime)
+
+	var finishAndGetRecording func() tracingpb.Recording
+	if s.onRecording != nil {
+		ctx, finishAndGetRecording = tracing.ContextWithRecordingSpan(ctx, s.Tracer, op)
+	}
+	f(ctx)
+	if finishAndGetRecording != nil {
+		s.onRecording(storeID, atDuration, finishAndGetRecording())
+	}
+}
+
 // tickQueues iterates over the next replicas for each store to
 // consider. It then enqueues each of these and ticks the replicate queue for
 // processing.
@@ -334,10 +353,16 @@ func (s *Simulator) tickQueues(ctx context.Context, tick time.Time, state state.
 
 		// Tick the split queue.
 		s.sqs[storeID].Tick(ctx, tick, state)
+
 		// Tick the replicate queue.
-		s.rqs[storeID].Tick(ctx, tick, state)
+		s.doAndMaybeTrace(ctx, storeID, tick, "replicateQueue.PlanOneChange", func(ctx context.Context) {
+			s.rqs[storeID].Tick(ctx, tick, state)
+		})
+
 		// Tick the lease queue.
-		s.lqs[storeID].Tick(ctx, tick, state)
+		s.doAndMaybeTrace(ctx, storeID, tick, "leaseQueue.PlanOneChange", func(ctx context.Context) {
+			s.lqs[storeID].Tick(ctx, tick, state)
+		})
 
 		// Tick changes that may have been enqueued with a lower completion
 		// than the current tick, from the queues.
@@ -389,14 +414,9 @@ 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())
-		}
+		s.doAndMaybeTrace(ctx, store.StoreID(), tick, "mma.ComputeChanges", func(ctx context.Context) {
+			s.mmSRs[store.StoreID()].Tick(ctx, tick, state)
+		})
 	}
 }
 

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

@@ -119,7 +119,7 @@ type SimulationSettings struct {
 	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)
+	OnRecording func(storeID int64, atDuration time.Duration, rec tracingpb.Recording)
 }
 
 // DefaultSimulationSettings returns a set of default settings for simulation.

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

@@ -8,6 +8,7 @@ go_library(
         "rand_framework.go",
         "rand_gen.go",
         "rand_util.go",
+        "tracer_helper.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/tests",
     visibility = ["//visibility:public"],
@@ -25,6 +26,9 @@ go_library(
         "//pkg/sql/catalog/catpb",
         "//pkg/sql/catalog/descpb",
         "//pkg/sql/catalog/multiregion",
+        "//pkg/util/tracing/tracingpb",
+        "@com_github_stretchr_testify//assert",
+        "@com_github_stretchr_testify//require",
     ],
 )
 
@@ -62,7 +66,6 @@ go_test(
         "//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

@@ -12,7 +12,6 @@ import (
 	"math/rand"
 	"os"
 	"path/filepath"
-	"regexp"
 	"strings"
 	"testing"
 	"time"
@@ -38,7 +37,6 @@ import (
 	"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"
 )
 
@@ -556,25 +554,9 @@ func TestDataDriven(t *testing.T) {
 
 						seedGen := rand.New(rand.NewSource(seed))
 						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
+							tr := makeTraceHelper(rewrite, plotDir, testName, sample+1, duration)
+							settingsGen.Settings.OnRecording = func(storeID int64, atDuration time.Duration, rec tracingpb.Recording) {
+								tr.OnRecording(t, storeID, atDuration, rec)
 							}
 
 							assertionFailures := []string{}

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

@@ -0,0 +1,71 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package tests
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type traceHelper struct {
+	enabled  bool
+	testName string
+	plotDir  string
+	sample   int
+
+	traceDirCreated  map[int64]struct{} // by storeID
+	atDurationFormat string
+}
+
+func (tr traceHelper) OnRecording(
+	t *testing.T, storeID int64, atDuration time.Duration, rec tracingpb.Recording,
+) {
+	if !tr.enabled || len(rec[0].Logs) == 0 {
+		return
+	}
+
+	traceDir := filepath.Join(tr.plotDir, "traces", fmt.Sprintf("s%d", storeID))
+	if _, ok := tr.traceDirCreated[storeID]; !ok {
+		tr.traceDirCreated[storeID] = struct{}{}
+		require.NoError(t, os.MkdirAll(traceDir, 0755))
+	}
+	re := regexp.MustCompile(`[^a-zA-Z0-9.]+`)
+
+	var sampleS string
+	if tr.sample > 0 {
+		sampleS = strconv.Itoa(tr.sample)
+	}
+	outName := fmt.Sprintf("%s%s_%s_%s_s%d", tr.testName, sampleS, fmt.Sprintf(tr.atDurationFormat, atDuration.Seconds()),
+		re.ReplaceAllString(rec[0].Operation, "_"), storeID)
+	assert.NoError(t, os.WriteFile(
+		filepath.Join(traceDir, outName),
+		[]byte(rec.String()), 0644))
+}
+
+func makeTraceHelper(
+	enabled bool, plotDir string, testName string, sample int, duration time.Duration,
+) traceHelper {
+	secondsSinceBeginningWidth := len(fmt.Sprintf("%d", int64(duration.Seconds()+1)))
+	return traceHelper{
+		enabled:         enabled,
+		testName:        testName,
+		plotDir:         plotDir,
+		sample:          sample,
+		traceDirCreated: map[int64]struct{}{},
+		// Print seconds with two decimal places. The decimal point and two decimal
+		// places need three more characters.
+		atDurationFormat: fmt.Sprintf("%%0%d.2fs", secondsSinceBeginningWidth+3),
+	}
+}