Merge #148444 #148697 #148764

148444: roachtest: global monitor expects specific process deaths r=herkolategan,srosenberg a=DarrylWong

This change refactors the global test monitor to expect specific process deaths rather than a number of any process deaths.  This change is to support the failure injection framework (#138958) as well as  failure injection in mixed version (#148084). In the future, we will want to be able to tell the state of the cluster/nodes from the monitor.

This change also refactors the mixed version tests to use the global monitor, as well as deprecates the old cluster monitor. See individual commits for details.

Informs: #118214
Release note: none

148697: allocator: hoist some calls to Locality r=wenyihu6 a=stevendanna

store.Locality() allocates. Previously we were calling this inside an double-nested loop in diversityRebalanceScore which itself is called in a loop.

This is probably just a drop in the bucket given how much is going on in the callers of this function, but perhaps it helps a bit:

    BenchmarkRebalanceToDiversityScore        427075              2734 ns/op            1920 B/op         12 allocs/op
    BenchmarkRebalanceToDiversityScore        904476              1443 ns/op             160 B/op          1 allocs/op

Informs #147800

Release note: None

148764: workload/schemachanger: fix invalid function syntax errors r=fqazi a=fqazi

Previously, the random schema changer workload used $$ as a delimiter for function definitions. Unfortunately, randomly generated strings used by functions could easily contain $$ causing the definition to terminator to show up. To address this patch modifies the delimiter to be `$FUNC_BODY$`.

Fixes: #148555

Release note: None

Co-authored-by: DarrylWong <[email protected]>
Co-authored-by: Steven Danna <[email protected]>
Co-authored-by: Faizan Qazi <[email protected]>

Author: 4 people

pkg/cmd/roachtest/cluster.go

@@ -2221,7 +2221,10 @@ func (c *clusterImpl) StartE(
 			return errors.Wrap(err, "failed to wait for replication after starting cockroach")
 		}
 	}
-
+	// If starting the cluster was successful, mark the nodes as healthy. N.B. we must wait
+	// until cluster startup succeeds as we may have tests that purposely inject failures into
+	// cluster startup.
+	c.t.Monitor().ExpectProcessAlive(nodes)
 	return nil
 }
 
@@ -2264,6 +2267,15 @@ func (c *clusterImpl) StartServiceForVirtualClusterE(
 			return err
 		}
 	}
+
+	// If we are starting a separate process virtual cluster, we need to
+	// mark each SQL instance as healthy.
+	if len(startOpts.SeparateProcessNodes) > 0 {
+		nodes := startOpts.SeparateProcessNodes
+		virtualClusterName := startOpts.RoachprodOpts.VirtualClusterName
+		sqlInstance := startOpts.RoachprodOpts.SQLInstance
+		c.t.Monitor().ExpectProcessAlive(nodes, option.VirtualClusterName(virtualClusterName), option.SQLInstance(sqlInstance))
+	}
 	return nil
 }
 
@@ -2290,6 +2302,9 @@ func (c *clusterImpl) StopServiceForVirtualClusterE(
 	nodes := c.All()
 	if len(stopOpts.SeparateProcessNodes) > 0 {
 		nodes = stopOpts.SeparateProcessNodes
+		virtualClusterName := stopOpts.RoachprodOpts.VirtualClusterName
+		sqlInstance := stopOpts.RoachprodOpts.SQLInstance
+		c.t.Monitor().ExpectProcessDead(nodes, option.VirtualClusterName(virtualClusterName), option.SQLInstance(sqlInstance))
 	}
 
 	return roachprod.StopServiceForVirtualCluster(
@@ -2396,6 +2411,7 @@ func (c *clusterImpl) StopE(
 		stopOpts.RoachprodOpts.Wait = true
 		stopOpts.RoachprodOpts.GracePeriod = 10
 	}
+	c.t.Monitor().ExpectProcessDead(selectedNodesOrDefault(nodes, c.All()))
 	return errors.Wrap(roachprod.Stop(ctx, l, c.MakeNodes(nodes...), stopOpts.RoachprodOpts), "cluster.StopE")
 }
 
@@ -2423,6 +2439,7 @@ func (c *clusterImpl) SignalE(
 	if c.spec.NodeCount == 0 {
 		return nil // unit tests
 	}
+	c.t.Monitor().ExpectProcessDead(selectedNodesOrDefault(nodes, c.All()))
 	return errors.Wrap(roachprod.Signal(ctx, l, c.MakeNodes(nodes...), sig), "cluster.Signal")
 }
 
@@ -2454,6 +2471,7 @@ func (c *clusterImpl) WipeE(
 	}
 	c.setStatusForClusterOpt("wiping", false, nodes...)
 	defer c.clearStatusForClusterOpt(false)
+	c.t.Monitor().ExpectProcessDead(selectedNodesOrDefault(nodes, c.All()))
 	return roachprod.Wipe(ctx, l, c.MakeNodes(nodes...), c.IsSecure())
 }
 
@@ -3009,14 +3027,16 @@ func (c *clusterImpl) Extend(ctx context.Context, d time.Duration, l *logger.Log
 	return nil
 }
 
-// NewMonitor creates a monitor that can watch for unexpected crdb node deaths on m.Wait()
+// NewDeprecatedMonitor creates a monitor that can watch for unexpected crdb node deaths on m.Wait()
 // and provide roachtest safe goroutines.
 //
 // As a general rule, if the user has a workload node, do not monitor it. A
 // monitor's semantics around handling expected node deaths breaks down if it's
 // monitoring a workload node.
-func (c *clusterImpl) NewMonitor(ctx context.Context, opts ...option.Option) cluster.Monitor {
-	return newMonitor(ctx, c.t, c, opts...)
+func (c *clusterImpl) NewDeprecatedMonitor(
+	ctx context.Context, opts ...option.Option,
+) cluster.Monitor {
+	return newMonitor(ctx, c.t, c, false /* expectExactProcessDeath */, opts...)
 }
 
 func (c *clusterImpl) StartGrafana(

pkg/cmd/roachtest/cluster/cluster_interface.go

@@ -66,7 +66,11 @@ type Cluster interface {
 	Stop(ctx context.Context, l *logger.Logger, stopOpts option.StopOpts, opts ...option.Option)
 	SignalE(ctx context.Context, l *logger.Logger, sig int, opts ...option.Option) error
 	Signal(ctx context.Context, l *logger.Logger, sig int, opts ...option.Option)
-	NewMonitor(context.Context, ...option.Option) Monitor
+
+	// NewDeprecatedMonitor is deprecated: See #118214
+	// Instead, use task.Tasker for goroutine management (e.g. test.Go) and test.Monitor
+	// for monitoring unexpected process deaths (e.g. through the test spec Monitor option)
+	NewDeprecatedMonitor(context.Context, ...option.Option) Monitor
 
 	// Starting virtual clusters.
 

pkg/cmd/roachtest/clusterstats/mock_cluster_generated_test.go

@@ -17,10 +17,66 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachprod"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/install"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/errors"
 	"golang.org/x/sync/errgroup"
 )
 
+// monitorProcess represents a single process that the monitor monitors.
+type monitorProcess struct {
+	node               install.Node
+	virtualClusterName string
+	sqlInstance        int
+}
+
+// MonitorExpectedProcessHealth represents the expected health of a process.
+type MonitorExpectedProcessHealth string
+
+const (
+	ExpectedAlive = MonitorExpectedProcessHealth("process alive")
+	ExpectedDead  = MonitorExpectedProcessHealth("process dead")
+)
+
+// expectedProcessHealth is a concurrent map that stores the expected health of
+// each registered monitorProcess. It is a thin wrapper over syncutil.Map that ensures
+// consistent naming of the system interface.
+type expectedProcessHealth struct {
+	syncutil.Map[monitorProcess, MonitorExpectedProcessHealth]
+}
+
+func newProcess(node install.Node, virtualClusterName string, sqlInstance int) monitorProcess {
+	if virtualClusterName == "" {
+		virtualClusterName = install.SystemInterfaceName
+	}
+	return monitorProcess{
+		node:               node,
+		virtualClusterName: virtualClusterName,
+		sqlInstance:        sqlInstance,
+	}
+}
+
+func (m *expectedProcessHealth) get(
+	node install.Node, virtualClusterName string, sqlInstance int,
+) MonitorExpectedProcessHealth {
+	val, ok := m.Load(newProcess(node, virtualClusterName, sqlInstance))
+	if !ok {
+		// If the process has no expected state, assume it should be healthy.
+		return ExpectedAlive
+	}
+	return *val
+}
+
+func (m *expectedProcessHealth) set(
+	nodes install.Nodes,
+	virtualClusterName string,
+	sqlInstance int,
+	health MonitorExpectedProcessHealth,
+) {
+	for _, node := range nodes {
+		m.Store(newProcess(node, virtualClusterName, sqlInstance), &health)
+	}
+}
+
 // monitorImpl implements the Monitor interface. A monitor both
 // manages "user tasks" -- goroutines provided by tests -- as well as
 // checks that every node in the cluster is still running. A failure
@@ -46,7 +102,16 @@ type monitorImpl struct {
 	monitorGroup *errgroup.Group // monitor goroutine
 	monitorOnce  sync.Once       // guarantees monitor goroutine is only started once
 
-	expDeaths int32 // atomically
+	// expExactProcessDeath if true indicates that the monitor should expect that a
+	// specified process, as denoted by the triple in expProcessHealth.get, is dead.
+	// Otherwise, the monitor will expect that only a certain number of process deaths.
+	// The former is a stronger assertion used in the new global roachtest monitor,
+	// while the latter should be removed when the deprecated cluster monitor is removed.
+	expExactProcessDeath bool
+	// Deprecated: This field is used by the deprecated cluster monitor to track the number
+	// of expected process deaths, and should be removed when the cluster monitor is removed.
+	expDeaths        int32 // atomically
+	expProcessHealth expectedProcessHealth
 }
 
 func newMonitor(
@@ -58,19 +123,46 @@ func newMonitor(
 		L() *logger.Logger
 	},
 	c cluster.Cluster,
+	expectExactProcessDeath bool,
 	opts ...option.Option,
 ) *monitorImpl {
 	m := &monitorImpl{
-		t:     t,
-		l:     t.L(),
-		nodes: c.MakeNodes(opts...),
+		t:                    t,
+		l:                    t.L(),
+		nodes:                c.MakeNodes(opts...),
+		expExactProcessDeath: expectExactProcessDeath,
+		expProcessHealth:     expectedProcessHealth{},
 	}
 	m.ctx, m.cancel = context.WithCancel(ctx)
 	m.userGroup, _ = errgroup.WithContext(m.ctx)
 	m.monitorGroup, _ = errgroup.WithContext(m.ctx)
 	return m
 }
 
+func (m *monitorImpl) ExpectProcessHealth(
+	nodes install.Nodes, health MonitorExpectedProcessHealth, opts ...option.OptionFunc,
+) {
+	var virtualClusterOptions option.VirtualClusterOptions
+	if err := option.Apply(&virtualClusterOptions, opts...); err != nil {
+		m.t.Fatal(err)
+	}
+	m.expProcessHealth.set(nodes, virtualClusterOptions.VirtualClusterName, virtualClusterOptions.SQLInstance, health)
+}
+
+// ExpectProcessDeath lets the monitor know that a set of processes are about
+// to be killed, and that their deaths should be ignored. Virtual cluster
+// options can be passed to denote a separate process.
+func (m *monitorImpl) ExpectProcessDead(nodes option.NodeListOption, opts ...option.OptionFunc) {
+	m.ExpectProcessHealth(nodes.InstallNodes(), ExpectedDead, opts...)
+}
+
+// ExpectProcessAlive lets the monitor know that a set of processes are
+// expected to be healthy. Virtual cluster options can be passed to denote
+// a separate process.
+func (m *monitorImpl) ExpectProcessAlive(nodes option.NodeListOption, opts ...option.OptionFunc) {
+	m.ExpectProcessHealth(nodes.InstallNodes(), ExpectedAlive, opts...)
+}
+
 // ExpectDeath lets the monitor know that a node is about to be killed, and that
 // this should be ignored.
 func (m *monitorImpl) ExpectDeath() {
@@ -190,12 +282,16 @@ func (m *monitorImpl) startNodeMonitor() {
 						)
 					}
 				case install.MonitorProcessDead:
-					isExpectedDeath := atomic.AddInt32(&m.expDeaths, -1) >= 0
-					if isExpectedDeath {
-						expectedDeathStr = ": expected"
+					var isExpectedDeath bool
+					if m.expExactProcessDeath {
+						isExpectedDeath = m.expProcessHealth.get(info.Node, e.VirtualClusterName, e.SQLInstance) == ExpectedDead
+					} else {
+						isExpectedDeath = atomic.AddInt32(&m.expDeaths, -1) >= 0
 					}
 
-					if !isExpectedDeath {
+					if isExpectedDeath {
+						expectedDeathStr = ": expected"
+					} else {
 						retErr = fmt.Errorf("unexpected node event: %s", info)
 					}
 				}

pkg/cmd/roachtest/monitor.go

@@ -191,7 +191,7 @@ type TestSpec struct {
 	Randomized bool
 
 	// Monitor specifies whether the test initiates a process monitor. Eventually,
-	// this should replace all instances of `cluster.NewMonitor`. To make this
+	// this should replace all instances of `cluster.NewDeprecatedMonitor`. To make this
 	// transition, tests should be modified to utilize the `test.Monitor` and
 	// `roachtestutil.Task` interfaces provided with each test.
 	Monitor bool

pkg/cmd/roachtest/registry/test_spec.go

@@ -283,20 +283,6 @@ func (h *Helper) GoCommand(cmd string, nodes option.NodeListOption) context.Canc
 	}, task.Name(desc))
 }
 
-// ExpectDeath alerts the testing infrastructure that a node is
-// expected to die. Regular restarts as part of the mixedversion
-// testing are already taken into account. This function should only
-// be used by tests that perform their own node restarts or chaos
-// events.
-func (h *Helper) ExpectDeath() {
-	h.ExpectDeaths(1)
-}
-
-// ExpectDeaths is the general version of `ExpectDeath()`.
-func (h *Helper) ExpectDeaths(n int) {
-	h.runner.monitor.ExpectDeaths(n)
-}
-
 // ClusterVersion returns the currently active cluster version. Avoids
 // querying the database if we are not running migrations, since the
 // test runner has cached version of the cluster versions.

pkg/cmd/roachtest/roachtestutil/mixedversion/helper.go

@@ -642,6 +642,10 @@ func NewTest(
 	crdbNodes option.NodeListOption,
 	options ...CustomOption,
 ) *Test {
+	if !t.Spec().(*registry.TestSpec).Monitor {
+		t.Fatal("mixedversion tests require enabling the global test monitor in the test spec")
+	}
+
 	opts := defaultTestOptions()
 	for _, fn := range options {
 		fn(&opts)
@@ -858,7 +862,7 @@ func (t *Test) Run() {
 }
 
 func (t *Test) run(plan *TestPlan) error {
-	return newTestRunner(t.ctx, t.cancel, plan, t.options.tag, t.logger, t.cluster).run()
+	return newTestRunner(t.ctx, t.cancel, plan, t.rt, t.options.tag, t.logger, t.cluster).run()
 }
 
 func (t *Test) plan() (plan *TestPlan, retErr error) {