Merge #140604 #141044 #141045 #141146

140604: roachtest: use atomic pointer for logger r=srosenberg,herkolategan a=DarrylWong

The test runner swaps out the test logger when running post test artifacts collection and checks. However, in the case of a timeout, the test goroutine may still be running and have access to the test logger. This leads to a race condition where the logger is replaced as it's being used by the test.

This change switches the test logger to use an atomic pointer instead.

Fixes: none
Epic: none
Release note: none

141044: physicalplan: minor fixes to pooling of FlowSpecs r=yuzefovich a=yuzefovich

**physicalplan: ensure that FlowSpecs are released**

We pool `FlowSpec` allocations in `GenerateFlowSpecs`, but previously we would only release them back to the pool on the main query path. This commit fixes that oversight.

**physicalplan: fix minor leak with pooling of FlowSpecs**

The only thing that we reuse when pooling `FlowSpec` objects is the capacity of `Processors` slice. Previously, we forgot to unset each element of that slice, so we could have hold onto to some processor specs which in turn could have hold onto some large objects (like `roachpb.Span`s in the TableReader) until a particular index of the slice is overwritten. This oversight is now fixed.

Found while looking at #140326.
Epic: None
Release note: None

141045: sqlstats: mutex improvements r=xinhaoz a=xinhaoz

Previously, when attempting to get a stmt or txn entry in the sql stats containers, we pass a flag to the get* functions specifying whether or not to create the entry if it does not exist. This commit cleans up this interface by splitting the function up into a pure get or tryCreate functions.

Note that originally we were pursuing a double checked lock approach, however Golang's RWMutex does not scale well with the number of CPUs. Instead this patch just simplifies the code here a bit to make some incoming changes that will eventually move this stats recording step out of the execution path easier. Once that happens we can investigate further improvment or removal of these mutexes.

Epic: none
Release note: None
Part of: #141024

### sqlstats: convert RWMutex to Mutex

This commit changes the RWMutex on the sql stats
containers to the regular exclusive Mutex struct.
The motivation behind this change is that RWMutex
scales poorly with CPU count. See benchmarks below.

Epic: none
Part of: #140590

141146: sqlstats: delete unused iterators r=kyle-a-wong a=kyle-a-wong

deletes unused iterators in persistedsqlstats

Epic: None
Release note: None

Co-authored-by: DarrylWong <[email protected]>
Co-authored-by: Yahor Yuzefovich <[email protected]>
Co-authored-by: Xin Hao Zhang <[email protected]>
Co-authored-by: Kyle Wong <[email protected]>

Author: 5 people

pkg/ccl/changefeedccl/changefeed_dist.go

@@ -555,7 +555,8 @@ func makePlan(
 
 		// Log the plan diagram URL so that we don't have to rely on it being in system.job_info.
 		const maxLenDiagURL = 1 << 20 // 1 MiB
-		flowSpecs := p.GenerateFlowSpecs()
+		flowSpecs, cleanup := p.GenerateFlowSpecs()
+		defer cleanup(flowSpecs)
 		if _, diagURL, err := execinfrapb.GeneratePlanDiagramURL(
 			fmt.Sprintf("changefeed: %d", jobID),
 			flowSpecs,

pkg/cmd/roachtest/github_test.go

@@ -139,12 +139,12 @@ func TestCreatePostRequest(t *testing.T) {
 
 		ti := &testImpl{
 			spec:        testSpec,
-			l:           nilLogger(),
 			start:       time.Date(2023, time.July, 21, 16, 34, 3, 817, time.UTC),
 			end:         time.Date(2023, time.July, 21, 16, 42, 13, 137, time.UTC),
 			cockroach:   "cockroach",
 			cockroachEA: "cockroach-ea",
 		}
+		ti.ReplaceL(nilLogger())
 
 		testClusterImpl := &clusterImpl{spec: clusterSpec, arch: vm.ArchAMD64, name: "foo"}
 		vo := vm.DefaultCreateOpts()

pkg/cmd/roachtest/test_impl.go

@@ -14,6 +14,7 @@ import (
 	"regexp"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry"
@@ -77,7 +78,11 @@ type testImpl struct {
 	buildVersion *version.Version
 
 	// l is the logger that the test will use for its output.
-	l *logger.Logger
+	//
+	// N.B. We need to use an atomic pointer here since the test
+	// runner can swap the logger out when running post test assertions
+	// and artifacts collection.
+	l atomic.Pointer[logger.Logger]
 
 	// taskManager manages tasks (goroutines) for tests.
 	taskManager task.Manager
@@ -172,7 +177,7 @@ func (t *testImpl) Cockroach() string {
 			// If the test is a benchmark test, we don't want to enable assertions
 			// as it will slow down performance.
 			if t.spec.Benchmark {
-				t.l.Printf("Benchmark test, running with standard cockroach")
+				t.L().Printf("Benchmark test, running with standard cockroach")
 				t.randomizedCockroach = t.StandardCockroach()
 				return
 			}
@@ -181,20 +186,20 @@ func (t *testImpl) Cockroach() string {
 				// The build with runtime assertions should exist in every nightly
 				// CI build, but we can't assume it exists in every roachtest call.
 				if path := t.RuntimeAssertionsCockroach(); path != "" {
-					t.l.Printf("Runtime assertions enabled")
+					t.L().Printf("Runtime assertions enabled")
 					t.randomizedCockroach = path
 					return
 				} else {
-					t.l.Printf("WARNING: running without runtime assertions since the corresponding binary was not specified")
+					t.L().Printf("WARNING: running without runtime assertions since the corresponding binary was not specified")
 				}
 			}
-			t.l.Printf("Runtime assertions disabled")
+			t.L().Printf("Runtime assertions disabled")
 			t.randomizedCockroach = t.StandardCockroach()
 		case registry.StandardCockroach:
-			t.l.Printf("Runtime assertions disabled: registry.StandardCockroach set")
+			t.L().Printf("Runtime assertions disabled: registry.StandardCockroach set")
 			t.randomizedCockroach = t.StandardCockroach()
 		case registry.RuntimeAssertionsCockroach:
-			t.l.Printf("Runtime assertions enabled: registry.RuntimeAssertionsCockroach set")
+			t.L().Printf("Runtime assertions enabled: registry.RuntimeAssertionsCockroach set")
 			t.randomizedCockroach = t.RuntimeAssertionsCockroach()
 		default:
 			t.Fatal("Specified cockroach binary does not exist.")
@@ -259,13 +264,12 @@ func (t *testImpl) SnapshotPrefix() string {
 
 // L returns the test's logger.
 func (t *testImpl) L() *logger.Logger {
-	return t.l
+	return t.l.Load()
 }
 
 // ReplaceL replaces the test's logger.
 func (t *testImpl) ReplaceL(l *logger.Logger) {
-	// TODO(tbg): get rid of this, this is racy & hacky.
-	t.l = l
+	t.l.Store(l)
 }
 
 func (t *testImpl) status(ctx context.Context, id int64, args ...interface{}) {

pkg/cmd/roachtest/test_impl_test.go

@@ -179,9 +179,8 @@ func Test_failuresMatchingError(t *testing.T) {
 }
 
 func Test_failureSpecifyOwnerAndAddFailureCombination(t *testing.T) {
-	ti := testImpl{
-		l: nilLogger(),
-	}
+	ti := testImpl{}
+	ti.ReplaceL(nilLogger())
 	ti.addFailure(0, "", vmPreemptionError("my_VM"))
 	errWithOwnership := failuresAsErrorWithOwnership(ti.failures())
 

pkg/cmd/roachtest/test_runner.go

@@ -830,14 +830,14 @@ func (r *testRunner) runWorker(
 			buildVersion:           binaryVersion,
 			artifactsDir:           testArtifactsDir,
 			artifactsSpec:          artifactsSpec,
-			l:                      testL,
 			versionsBinaryOverride: topt.versionsBinaryOverride,
 			skipInit:               topt.skipInit,
 			debug:                  clustersOpt.debugMode.IsDebug(),
 			goCoverEnabled:         topt.goCoverEnabled,
 			exportOpenmetrics:      topt.exportOpenMetrics,
 			runID:                  generateRunID(clustersOpt),
 		}
+		t.ReplaceL(testL)
 		github := newGithubIssues(r.config.disableIssue, c, vmCreateOpts)
 
 		// handleClusterCreationFailure can be called when the `err` given

pkg/cmd/roachtest/test_test.go

@@ -831,3 +831,38 @@ func TestVMPreemptionPolling(t *testing.T) {
 		require.NoError(t, err)
 	})
 }
+
+// TestRunnerFailureAfterTimeout checks that a test has a failure added
+// after the test has timed out works as expected.
+//
+// Specifically, this is a regression test that replacing the test logger
+// for post test artifacts collection or assertion checks is atomic and
+// doesn't race with the logger potentially still being used by the test.
+func TestRunnerFailureAfterTimeout(t *testing.T) {
+	ctx := context.Background()
+	stopper := stop.NewStopper()
+	defer stopper.Stop(ctx)
+	cr := newClusterRegistry()
+	runner := newUnitTestRunner(cr, stopper)
+
+	var buf syncedBuffer
+	copt := defaultClusterOpt()
+	lopt := defaultLoggingOpt(&buf)
+	test := registry.TestSpec{
+		Name:  `timeout`,
+		Owner: OwnerUnitTest,
+		// Set the timeout very low so we can observe the timeout
+		// and error racing.
+		Timeout:          1 * time.Nanosecond,
+		Cluster:          spec.MakeClusterSpec(0),
+		CompatibleClouds: registry.AllExceptAWS,
+		Suites:           registry.Suites(registry.Nightly),
+		CockroachBinary:  registry.StandardCockroach,
+		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
+			t.Error("test failed")
+		},
+	}
+	err := runner.Run(ctx, []registry.TestSpec{test}, 1, /* count */
+		defaultParallelism, copt, testOpts{}, lopt)
+	require.Error(t, err)
+}

pkg/jobs/jobsprofiler/profiler.go

@@ -41,7 +41,8 @@ func StorePlanDiagram(
 		ctx, cancel = stopper.WithCancelOnQuiesce(ctx)
 		defer cancel()
 
-		flowSpecs := p.GenerateFlowSpecs()
+		flowSpecs, cleanup := p.GenerateFlowSpecs()
+		defer cleanup(flowSpecs)
 		_, diagURL, err := execinfrapb.GeneratePlanDiagramURL(fmt.Sprintf("job:%d", jobID), flowSpecs,
 			execinfrapb.DiagramFlags{})
 		if err != nil {

pkg/sql/distsql_plan_bulk.go

@@ -121,8 +121,10 @@ type PhysicalPlanMaker func(context.Context, *DistSQLPlanner) (*PhysicalPlan, *P
 // the number in the old one.
 func calculatePlanGrowth(before, after *PhysicalPlan) (int, float64) {
 	var changed int
-	beforeSpecs := before.GenerateFlowSpecs()
-	afterSpecs := after.GenerateFlowSpecs()
+	beforeSpecs, beforeCleanup := before.GenerateFlowSpecs()
+	defer beforeCleanup(beforeSpecs)
+	afterSpecs, afterCleanup := after.GenerateFlowSpecs()
+	defer afterCleanup(afterSpecs)
 
 	// How many nodes are in after that are not in before, or are in both but
 	// have changed their spec?

pkg/sql/distsql_running.go

@@ -675,7 +675,8 @@ func (dsp *DistSQLPlanner) Run(
 	evalCtx *extendedEvalContext,
 	finishedSetupFn func(localFlow flowinfra.Flow),
 ) {
-	flows := plan.GenerateFlowSpecs()
+	// Ignore the cleanup function since we will release each spec separately.
+	flows, _ := plan.GenerateFlowSpecs()
 	gatewayFlowSpec, ok := flows[dsp.gatewaySQLInstanceID]
 	if !ok {
 		recv.SetError(errors.Errorf("expected to find gateway flow"))

pkg/sql/explain_plan.go

@@ -88,7 +88,8 @@ func (e *explainPlanNode) startExec(params runParams) error {
 			// cause an error or panic, so swallow the error. See #40677 for example.
 			finalizePlanWithRowCount(params.ctx, planCtx, physicalPlan, plan.mainRowCount)
 			ob.AddDistribution(physicalPlan.Distribution.String())
-			flows := physicalPlan.GenerateFlowSpecs()
+			flows, flowsCleanup := physicalPlan.GenerateFlowSpecs()
+			defer flowsCleanup(flows)
 
 			ctxSessionData := planCtx.EvalContext().SessionData()
 			var willVectorize bool

pkg/sql/explain_vec.go

@@ -56,7 +56,8 @@ func (n *explainVecNode) startExec(params runParams) error {
 	}
 
 	finalizePlanWithRowCount(params.ctx, planCtx, physPlan, n.plan.mainRowCount)
-	flows := physPlan.GenerateFlowSpecs()
+	flows, flowsCleanup := physPlan.GenerateFlowSpecs()
+	defer flowsCleanup(flows)
 	flowCtx := newFlowCtxForExplainPurposes(params.ctx, params.p)
 
 	// We want to get the vectorized plan which would be executed with the

pkg/sql/physicalplan/physical_plan.go

@@ -979,7 +979,14 @@ func (p *PhysicalPlan) PopulateEndpoints() {
 
 // GenerateFlowSpecs takes a plan (with populated endpoints) and generates the
 // set of FlowSpecs (one per node involved in the plan).
-func (p *PhysicalPlan) GenerateFlowSpecs() map[base.SQLInstanceID]*execinfrapb.FlowSpec {
+//
+// The returned function should be called whenever the caller is done with all
+// FlowSpecs. The caller is free to ignore it if the specs will be released
+// separately.
+func (p *PhysicalPlan) GenerateFlowSpecs() (
+	_ map[base.SQLInstanceID]*execinfrapb.FlowSpec,
+	cleanup func(map[base.SQLInstanceID]*execinfrapb.FlowSpec),
+) {
 	flowID := execinfrapb.FlowID{
 		UUID: p.FlowID,
 	}
@@ -993,7 +1000,15 @@ func (p *PhysicalPlan) GenerateFlowSpecs() map[base.SQLInstanceID]*execinfrapb.F
 		}
 		flowSpec.Processors = append(flowSpec.Processors, proc.Spec)
 	}
-	return flows
+	// Note that we don't return an anonymous function with no arguments to not
+	// incur an allocation.
+	return flows, releaseAll
+}
+
+func releaseAll(flows map[base.SQLInstanceID]*execinfrapb.FlowSpec) {
+	for _, flowSpec := range flows {
+		ReleaseFlowSpec(flowSpec)
+	}
 }
 
 // SetRowEstimates updates p according to the row estimates of left and right

pkg/sql/physicalplan/specs.go

@@ -34,6 +34,7 @@ func ReleaseFlowSpec(spec *execinfrapb.FlowSpec) {
 		if tr := spec.Processors[i].Core.TableReader; tr != nil {
 			releaseTableReaderSpec(tr)
 		}
+		spec.Processors[i] = execinfrapb.ProcessorSpec{}
 	}
 	*spec = execinfrapb.FlowSpec{
 		Processors: spec.Processors[:0],

pkg/sql/sqlstats/persistedsqlstats/BUILD.bazel

@@ -6,16 +6,12 @@ go_library(
     name = "persistedsqlstats",
     srcs = [
         "cluster_settings.go",
-        "combined_iterator.go",
         "compaction_exec.go",
         "compaction_scheduling.go",
         "controller.go",
         "flush.go",
-        "mem_iterator.go",
         "provider.go",
         "scheduled_job_monitor.go",
-        "stmt_reader.go",
-        "txn_reader.go",
     ],
     embed = [":persistedsqlstats_go_proto"],
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/sqlstats/persistedsqlstats",
@@ -40,7 +36,6 @@ go_library(
         "//pkg/sql/sqlstats/persistedsqlstats/sqlstatsutil",
         "//pkg/sql/sqlstats/sslocal",
         "//pkg/sql/types",
-        "//pkg/util",
         "//pkg/util/admission/admissionpb",
         "//pkg/util/log",
         "//pkg/util/metric",
@@ -66,7 +61,6 @@ go_test(
         "datadriven_test.go",
         "flush_test.go",
         "main_test.go",
-        "reader_test.go",
         "scheduled_sql_stats_compaction_test.go",
     ],
     data = glob(["testdata/**"]),