Merge cockroachdb#157132 cockroachdb#157133 cockroachdb#157150

157132: tracing: add minimal format suitable for datadriven tests r=wenyihu6 a=tbg

I'd like to make the trace events emitted from deterministic multi-metric allocator tests part of the expected test outputs.

I realized there wasn't a standardized deterministic and barebones output format, so I added one. In the test I have in mind I'll probably condense the file:line prefix further, but it seemed unsavory to slap regexp transformations into the tracing packages, and unfortunately that is required to solve the problem if we want to avoid changing the `log` package.

Related to cockroachdb#157133.

Epic: none

157133: mmaprototype: add tracing to testing for rebalanceStores r=tbg a=tbg

Stacked on cockroachdb#157132. Don't review tracing pkg code here.

----

This PR adds the basic capability to exercise rebalanceStores through
TestClusterState. It does this by moving rebalanceStores from allocatorState to
clusterState, which is a small and mostly mechanical refactor.

The immediate motivation behind this change is being able to exercise
the functionality for cockroachdb#156776 while it is being added, while being able
to see the internal tracepoints as well.

Informs cockroachdb#156776.

Epic: CRDB-55052


157150: cloud/azure: add optional azure client caching r=dt a=dt

This adds a signle-slot, process-wide in-memory cache of the most recently used azure SDK storage client allowing subsequent attempts to open clients with the same configuration to reuse the already open client. This can allow the SDK's internal caching of things like IMDS tokens to be used across separate CRDB operations on separate files which currently each open their own client.

This can be particularly impactful for avoiding hitting the IMDS rate limit.

Release note: none.
Epic: none.

Co-authored-by: Tobias Grieger <[email protected]>
Co-authored-by: David Taylor <[email protected]>

Author: 3 people

pkg/BUILD.bazel

@@ -835,6 +835,7 @@ ALL_TESTS = [
     "//pkg/util/tracing/goexectrace:goexectrace_test",
     "//pkg/util/tracing/grpcinterceptor:grpcinterceptor_test",
     "//pkg/util/tracing/service:service_test",
+    "//pkg/util/tracing/tracingpb:tracingpb_test",
     "//pkg/util/tracing:tracing_test",
     "//pkg/util/treeprinter:treeprinter_test",
     "//pkg/util/trigram:trigram_test",
@@ -2842,6 +2843,7 @@ GO_TARGETS = [
     "//pkg/util/tracing/service:service",
     "//pkg/util/tracing/service:service_test",
     "//pkg/util/tracing/tracingpb:tracingpb",
+    "//pkg/util/tracing/tracingpb:tracingpb_test",
     "//pkg/util/tracing/tracingservicepb:tracingservicepb",
     "//pkg/util/tracing/tracingui:tracingui",
     "//pkg/util/tracing/tracingutil:tracingutil",

pkg/cloud/azure/BUILD.bazel

@@ -24,6 +24,7 @@ go_library(
         "//pkg/util/envutil",
         "//pkg/util/ioctx",
         "//pkg/util/log",
+        "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "//pkg/util/tracing",
         "@com_github_azure_azure_sdk_for_go_sdk_azcore//:azcore",

pkg/cloud/azure/azure_storage.go

@@ -29,6 +29,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util/ioctx"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/errors"
@@ -55,6 +56,13 @@ var tryTimeout = settings.RegisterDurationSetting(
 	"the timeout for individual retry attempts in Azure operations",
 	60*time.Second)
 
+var reuseSession = settings.RegisterBoolSetting(
+	settings.ApplicationLevel,
+	"cloudstorage.azure.session_reuse.enabled",
+	"persist the last opened azure client and re-use it when opening a new client with the same argument (some settings may take 2mins to take effect)",
+	false,
+)
+
 // A note on Azure authentication:
 //
 // The standardized way to authenticate a third-party identity to the Azure
@@ -214,6 +222,14 @@ type azureStorage struct {
 	settings  *cluster.Settings
 }
 
+var azClientCache struct {
+	syncutil.Mutex
+	// TODO(dt): make this an >1 item cache e.g. add a FIFO ring.
+	key    cloudpb.ExternalStorage_Azure
+	set    time.Time
+	client *service.Client
+}
+
 var _ cloud.ExternalStorage = &azureStorage{}
 
 func makeAzureStorage(
@@ -224,6 +240,7 @@ func makeAzureStorage(
 	if conf == nil {
 		return nil, errors.Errorf("azure upload requested but info missing")
 	}
+
 	env, err := azure.EnvironmentFromName(conf.Environment)
 	if err != nil {
 		return nil, errors.Wrap(err, "azure environment")
@@ -256,6 +273,29 @@ func makeAzureStorage(
 	opts.Retry.TryTimeout = tryTimeout.Get(&args.Settings.SV)
 
 	var azClient *service.Client
+
+	clientConf := *conf
+	clientConf.Prefix = "" // Prefix is not part of the client identity.
+
+	if reuseSession.Get(&args.Settings.SV) {
+		func() {
+			azClientCache.Lock()
+			defer azClientCache.Unlock()
+			if cached := azClientCache.client; cached != nil && azClientCache.key == clientConf && timeutil.Since(azClientCache.set) < 2*time.Minute {
+				azClient = cached
+			}
+		}()
+		if azClient != nil {
+			return &azureStorage{
+				conf:      conf,
+				ioConf:    args.IOConf,
+				container: azClient.NewContainerClient(conf.Container),
+				prefix:    conf.Prefix,
+				settings:  args.Settings,
+			}, nil
+		}
+	}
+
 	switch conf.Auth {
 	case cloudpb.AzureAuth_LEGACY:
 		credential, err := azblob.NewSharedKeyCredential(conf.AccountName, conf.AccountKey)
@@ -299,6 +339,14 @@ func makeAzureStorage(
 		return nil, errors.Errorf("unsupported value %s for %s", conf.Auth, cloud.AuthParam)
 	}
 
+	if reuseSession.Get(&args.Settings.SV) {
+		azClientCache.Lock()
+		defer azClientCache.Unlock()
+		azClientCache.key = clientConf
+		azClientCache.client = azClient
+		azClientCache.set = timeutil.Now()
+	}
+
 	return &azureStorage{
 		conf:      conf,
 		ioConf:    args.IOConf,

pkg/kv/kvserver/allocator/mmaprototype/BUILD.bazel

@@ -53,8 +53,10 @@ go_test(
         "//pkg/testutils/dd",
         "//pkg/util/randutil",
         "//pkg/util/timeutil",
+        "//pkg/util/tracing",
         "@com_github_cockroachdb_datadriven//:datadriven",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_redact//:redact",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/kv/kvserver/allocator/mmaprototype/cluster_state_test.go

@@ -21,7 +21,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/dd"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/datadriven"
+	"github.com/cockroachdb/redact"
 	"github.com/stretchr/testify/require"
 )
 
@@ -519,8 +521,15 @@ func TestClusterState(t *testing.T) {
 					storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
 					rng := rand.New(rand.NewSource(0))
 					dsm := newDiversityScoringMemo()
-					cs.rebalanceStores(context.Background(), storeID, rng, dsm)
-					return printPendingChangesTest(testingGetPendingChanges(t, cs))
+					tr := tracing.NewTracer()
+					tr.SetRedactable(true)
+					defer tr.Close()
+					ctx, finishAndGet := tracing.ContextWithRecordingSpan(context.Background(), tr, "rebalance-stores")
+					cs.rebalanceStores(ctx, storeID, rng, dsm)
+					rec := finishAndGet()
+					var sb redact.StringBuilder
+					rec.SafeFormatMinimal(&sb)
+					return sb.String() + printPendingChangesTest(testingGetPendingChanges(t, cs))
 
 				case "tick":
 					seconds := dd.ScanArg[int](t, d, "seconds")

pkg/kv/kvserver/allocator/mmaprototype/testdata/cluster_state/rebalance_stores_basic

@@ -55,10 +55,15 @@ store-id=2 node-id=2 status=ok accepting all reported=[cpu:10, write-bandwidth:0
 store-id=3 node-id=3 status=ok accepting all reported=[cpu:10, write-bandwidth:0, byte-size:0] adjusted=[cpu:10, write-bandwidth:0, byte-size:0] node-reported-cpu=10 node-adjusted-cpu=10 seq=1
   top-k-ranges (local-store-id=1) dim=WriteBandwidth: r1
 
-# TODO(tbg): maybe it'll be good to at least optionally be able to have the
-# trace printed in the output here.
+# s1 is overloaded and local, so rebalanceStores should try and succeed to shed
+# a lease from it.
 rebalance-stores store-id=1
 ----
+[mmaid=1] cluster means: (stores-load [cpu:33, write-bandwidth:0, byte-size:0]) (stores-capacity [cpu:100, write-bandwidth:100, byte-size:100]) (nodes-cpu-load 33) (nodes-cpu-capacity 100)
+[mmaid=1] overload-continued s1 ((store=overloadUrgent worst=CPURate cpu=overloadUrgent writes=loadNormal bytes=loadNormal node=overloadUrgent high_disk=false frac_pending=0.00,0.00(true))) - within grace period
+[mmaid=1] start processing shedding store s1: cpu node load overloadUrgent, store load overloadUrgent, worst dim CPURate
+[mmaid=1] top-K[CPURate] ranges for s1 with lease on local s1: r1:[cpu:60, write-bandwidth:0, byte-size:0]
+[mmaid=1] result(success): shedding r1 lease from s1 to s2 [change:r1=[transfer_to=2 cids=1,2]] with resulting loads source:[cpu:40, write-bandwidth:0, byte-size:0] target:[cpu:54, write-bandwidth:0, byte-size:0] (means: [cpu:33, write-bandwidth:0, byte-size:0]) (frac_pending: (src:0.00,target:0.50) (src:4.40,target:0.00))
 pending(2)
 change-id=1 store-id=1 node-id=1 range-id=1 load-delta=[cpu:-40, write-bandwidth:0, byte-size:0] start=0s
   prev=(replica-id=1 type=VOTER_FULL leaseholder=true)

pkg/util/tracing/tracingpb/BUILD.bazel

@@ -1,6 +1,6 @@
 load("@rules_proto//proto:defs.bzl", "proto_library")
 load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "tracingpb",
@@ -51,3 +51,20 @@ go_proto_library(
         "@com_github_gogo_protobuf//gogoproto",
     ],
 )
+
+go_test(
+    name = "tracingpb_test",
+    srcs = [
+        "condense_test.go",
+        "recording_test.go",
+    ],
+    data = glob(["testdata/**"]),
+    embed = [":tracingpb"],
+    deps = [
+        "//pkg/testutils/datapathutils",
+        "//pkg/testutils/echotest",
+        "//pkg/util/log",
+        "//pkg/util/tracing",
+        "@com_github_cockroachdb_redact//:redact",
+    ],
+)

pkg/util/tracing/tracingpb/condense_test.go

@@ -0,0 +1,79 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package tracingpb
+
+import (
+	"testing"
+
+	"github.com/cockroachdb/redact"
+)
+
+func TestCondensePathLinePrefix(t *testing.T) {
+	for _, test := range []struct {
+		name     string
+		input    string
+		expected string
+	}{
+		{
+			name:     "no_directory",
+			input:    "foo.go:1234 message",
+			expected: "message",
+		},
+		{
+			name:     "full_path",
+			input:    "pkg/foo/bar.go:123 message",
+			expected: "message",
+		},
+		{
+			// Won't redact because there are redaction markers in the path.
+			// This doesn't happen in the real world but we want to make sure
+			// condensePathLinePrefix doesn't mess with redaction at all.
+			name:     "redaction_in_path",
+			input:    "pkg/‹foo›/bar.go:123 message",
+			expected: "pkg/‹foo›/bar.go:123 message",
+		},
+		{
+			// Redaction in message is expected and okay, since
+			// condensePathLinePrefix is written in a way to
+			// preserve that.
+			name:     "redaction_in_message",
+			input:    "pkg/foo/bar.go:123 ‹sensitive› message",
+			expected: "‹sensitive› message",
+		},
+		{
+			// Not expecting it in the real world since we only call
+			// condensePathLinePrefix in formatMinimal, but if we get formatFull, it
+			// also works.
+			name:     "full_format",
+			input:    "event: pkg/foo/bar.go:1234 ‹sensitive› message",
+			expected: "‹sensitive› message",
+		},
+		{
+			name:     "unredacted_trace",
+			input:    "‹pkg/foo/bar.go:1234 some message›",
+			expected: "‹some message›",
+		},
+		{
+			// Should never happen (there would be escaping). But either way, works
+			// just the same.
+			name:     "unredacted_trace_nested_unescaped",
+			input:    "‹pkg/foo/bar.go:1234 some ‹message››",
+			expected: "‹some ‹message››",
+		},
+		{
+			name:     "unredacted_trace_nested_escaped",
+			input:    "‹pkg/foo/bar.go:1234 " + string(redact.EscapeMarkers([]byte("some ‹message›"))) + "›",
+			expected: "‹some ?message?›", // NB: this is verbatim the escaped string inserted above
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			got := string(condensePathLinePrefix(redact.RedactableString(test.input)))
+			if got != test.expected {
+				t.Errorf("expected %q, got %q", test.expected, got)
+			}
+		})
+	}
+}