Merge #40229 #40271 #40320

40229: demo: Add location information to default demo localities. r=knz a=rohany

Fixes #39937.

Adds default latitude and longitude information for the default
localities in cockroach demo so that the web UI is pre-populated.

Release note (cli change): Add location information to default
localities to populate web UI for cockroach demo.

40271: demo: Add a Movr workload to run in cockroach demo r=danhhz,nstewart a=rohany

Fixes #39944.

When the `--with-load` flag is passed to demo, a sample movr
workload is run against the database.

I tried my best to replicate the Movr load in Nate's example script.

Based on talks with Dan, it was better to copy out a simpler
version of the logic to run the workload than try to hook
into workload.

Release note (cli change): Add a workload to run with cockroach demo.

40320: randutil: allow COCKROACH_RANDOM_SEED to seed NewPseudoRand r=mjibson a=mjibson

Multiple people have been confused about the lack of interaction
between NewPseudoRand (which used to always generate a random seed)
and COCKROACH_RANDOM_SEED (which only influenced using the global rand
methods).

Release note: None

Co-authored-by: Rohan Yadav <rohany@alumni.cmu.edu>
Co-authored-by: Matt Jibson <matt.jibson@gmail.com>

pkg/cli/cliflags/flags.go

@@ -897,6 +897,11 @@ The line length where sqlfmt will try to wrap.`,
 		Description: `How many in-memory nodes to create for the demo.`,
 	}
 
+	RunDemoWorkload = FlagInfo{
+		Name:        "with-load",
+		Description: `Run a demo workload against the pre-loaded database.`,
+	}
+
 	DemoNodeLocality = FlagInfo{
 		Name: "demo-locality",
 		Description: `

pkg/cli/context.go

@@ -145,6 +145,7 @@ func initCLIDefaults() {
 
 	demoCtx.nodes = 1
 	demoCtx.useEmptyDatabase = false
+	demoCtx.runWorkload = false
 	demoCtx.localities = nil
 
 	initPreFlagsDefaults()
@@ -337,5 +338,6 @@ var sqlfmtCtx struct {
 var demoCtx struct {
 	nodes            int
 	useEmptyDatabase bool
+	runWorkload      bool
 	localities       demoLocalityList
 }

pkg/cli/demo.go

@@ -15,6 +15,7 @@ import (
 	gosql "database/sql"
 	"fmt"
 	"net/url"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/cli/cliflags"
@@ -25,11 +26,14 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/log/logflags"
+	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/workload"
+	"github.com/cockroachdb/cockroach/pkg/workload/histogram"
 	"github.com/cockroachdb/cockroach/pkg/workload/workloadsql"
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
+	"golang.org/x/time/rate"
 )
 
 var demoCmd = &cobra.Command{
@@ -52,7 +56,8 @@ to avoid pre-loading a dataset.`,
 const defaultGeneratorName = "movr"
 
 var defaultGenerator workload.Generator
-var defaultLocalities = []roachpb.Locality{
+
+var defaultLocalities = demoLocalityList{
 	// Default localities for a 3 node cluster
 	{Tiers: []roachpb.Tier{{Key: "region", Value: "us-east1"}, {Key: "az", Value: "b"}}},
 	{Tiers: []roachpb.Tier{{Key: "region", Value: "us-east1"}, {Key: "az", Value: "c"}}},
@@ -206,17 +211,73 @@ func setupTransientServers(
 		if _, err := workloadsql.Setup(ctx, db, gen, l); err != nil {
 			return ``, ``, cleanup, err
 		}
+
+		if demoCtx.runWorkload {
+			if err := runWorkload(ctx, gen, urlStr, stopper); err != nil {
+				return ``, ``, cleanup, err
+			}
+		}
 	}
 
 	return urlStr, s.AdminURL(), cleanup, nil
 }
 
+func runWorkload(
+	ctx context.Context, gen workload.Generator, dbURL string, stopper *stop.Stopper,
+) error {
+	opser, ok := gen.(workload.Opser)
+	if !ok {
+		return errors.Errorf("default dataset %s does not have a workload defined", gen.Meta().Name)
+	}
+
+	// Dummy registry to prove to the Opser.
+	reg := histogram.NewRegistry(time.Duration(100) * time.Millisecond)
+	ops, err := opser.Ops([]string{dbURL}, reg)
+	if err != nil {
+		return errors.Wrap(err, "unable to create workload")
+	}
+
+	// Use a light rate limit of 25 queries per second
+	limiter := rate.NewLimiter(rate.Limit(25), 1)
+
+	// Start a goroutine to run each of the workload functions.
+	for _, workerFn := range ops.WorkerFns {
+		workloadFun := func(f func(context.Context) error) func(context.Context) {
+			return func(ctx context.Context) {
+				for {
+					// Limit how quickly we can generate work.
+					if err := limiter.Wait(ctx); err != nil {
+						// When the limiter throws an error, panic because we don't
+						// expect any errors from it.
+						panic(err)
+					}
+					if err := f(ctx); err != nil {
+						// Only log an error and return when the workload function throws
+						// an error, because errors these errors should be ignored, and
+						// should not interrupt the rest of the demo.
+						log.Warningf(ctx, "Error running workload query: %+v\n", err)
+						return
+					}
+				}
+			}
+		}
+		stopper.RunWorker(ctx, workloadFun(workerFn))
+	}
+
+	return nil
+}
+
 func runDemo(cmd *cobra.Command, gen workload.Generator) error {
 	if gen == nil && !demoCtx.useEmptyDatabase {
 		// Use a default dataset unless prevented by --empty.
 		gen = defaultGenerator
 	}
 
+	// Make sure that the user didn't request a workload and an empty database.
+	if demoCtx.runWorkload && demoCtx.useEmptyDatabase {
+		return errors.New("cannot run a workload against an empty database")
+	}
+
 	connURL, adminURL, cleanup, err := setupTransientServers(cmd, gen)
 	defer cleanup()
 	if err != nil {

pkg/cli/flags.go

@@ -587,6 +587,7 @@ func init() {
 	// The --empty flag is only valid for the top level demo command,
 	// so we use the regular flag set.
 	BoolFlag(demoCmd.Flags(), &demoCtx.useEmptyDatabase, cliflags.UseEmptyDatabase, false)
+	BoolFlag(demoCmd.Flags(), &demoCtx.runWorkload, cliflags.RunDemoWorkload, false)
 	VarFlag(demoFlags, &demoCtx.localities, cliflags.DemoNodeLocality)
 
 	// sqlfmt command.

pkg/internal/sqlsmith/sqlsmith_test.go

@@ -14,7 +14,6 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	"math/rand"
 	"reflect"
 	"strings"
 	"testing"
@@ -49,10 +48,7 @@ func TestGenerateParse(t *testing.T) {
 	s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{})
 	defer s.Stopper().Stop(ctx)
 
-	// Set COCKROACH_RANDOM_SEED to make this test deterministic between
-	// runs.
-	randutil.SeedForTests()
-	rnd := rand.New(rand.NewSource(rand.Int63()))
+	rnd, _ := randutil.NewPseudoRand()
 
 	db := sqlutils.MakeSQLRunner(sqlDB)
 	var opts []SmitherOption

pkg/roachpb/metadata.go

@@ -515,3 +515,37 @@ func (l *Locality) Set(value string) error {
 	l.Tiers = tiers
 	return nil
 }
+
+// DefaultLocationInformation is used to populate the system.locations
+// table. The region values here are specific to GCP.
+var DefaultLocationInformation = []struct {
+	Locality  Locality
+	Latitude  string
+	Longitude string
+}{
+	{
+		Locality:  Locality{Tiers: []Tier{{Key: "region", Value: "us-east1"}}},
+		Latitude:  "33.836082",
+		Longitude: "-81.163727",
+	},
+	{
+		Locality:  Locality{Tiers: []Tier{{Key: "region", Value: "us-east4"}}},
+		Latitude:  "37.478397",
+		Longitude: "-76.453077",
+	},
+	{
+		Locality:  Locality{Tiers: []Tier{{Key: "region", Value: "us-central1"}}},
+		Latitude:  "42.032974",
+		Longitude: "-93.581543",
+	},
+	{
+		Locality:  Locality{Tiers: []Tier{{Key: "region", Value: "us-west1"}}},
+		Latitude:  "43.804133",
+		Longitude: "-120.554201",
+	},
+	{
+		Locality:  Locality{Tiers: []Tier{{Key: "region", Value: "europe-west1"}}},
+		Latitude:  "50.44816",
+		Longitude: "3.81886",
+	},
+}

pkg/sqlmigrations/migrations.go

@@ -206,6 +206,11 @@ var backwardCompatibleMigrations = []migrationDescriptor{
 		name:   "propagate the ts purge interval to the new setting names",
 		workFn: retireOldTsPurgeIntervalSettings,
 	},
+	{
+		// Introduced in ? TODO (rohany): what version is this?
+		name:   "update system.locations with default location data",
+		workFn: updateSystemLocationData,
+	},
 }
 
 func staticIDs(ids ...sqlbase.ID) func(ctx context.Context, db db) ([]sqlbase.ID, error) {
@@ -918,3 +923,27 @@ ON CONFLICT (name) DO NOTHING`,
 
 	return nil
 }
+
+func updateSystemLocationData(ctx context.Context, r runner) error {
+	// See if the system.locations table already has data in it.
+	// If so, we don't want to do anything.
+	row, err := r.sqlExecutor.QueryRow(ctx, "update-system-locations",
+		nil, `SELECT count(*) FROM system.locations`)
+	if err != nil {
+		return err
+	}
+	count := int(tree.MustBeDInt(row[0]))
+	if count != 0 {
+		return nil
+	}
+
+	for _, loc := range roachpb.DefaultLocationInformation {
+		stmt := `UPSERT INTO system.locations VALUES ($1, $2, $3, $4)`
+		tier := loc.Locality.Tiers[0]
+		if _, err := r.sqlExecutor.Exec(ctx, "update-system-locations", nil,
+			stmt, tier.Key, tier.Value, loc.Latitude, loc.Longitude); err != nil {
+			return err
+		}
+	}
+	return nil
+}

pkg/sqlmigrations/migrations_test.go

@@ -639,3 +639,32 @@ func TestExpectedInitialRangeCount(t *testing.T) {
 		return nil
 	})
 }
+
+func TestUpdateSystemLocationData(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	ctx := context.Background()
+
+	mt := makeMigrationTest(ctx, t)
+	defer mt.close(ctx)
+
+	migration := mt.pop(t, "update system.locations with default location data")
+	mt.start(t, base.TestServerArgs{})
+
+	// Check that we don't have any data in the system.locations table without the migration.
+	var count int
+	mt.sqlDB.QueryRow(t, `SELECT count(*) FROM system.locations`).Scan(&count)
+	if count != 0 {
+		t.Fatalf("Exected to find 0 rows in system.locations. Found  %d instead", count)
+	}
+
+	// Run the migration to insert locations.
+	if err := mt.runMigration(ctx, migration); err != nil {
+		t.Errorf("expected success, got %q", err)
+	}
+
+	// Check that we have all of the expected locations.
+	mt.sqlDB.QueryRow(t, `SELECT count(*) FROM system.locations`).Scan(&count)
+	if count != len(roachpb.DefaultLocationInformation) {
+		t.Fatalf("Exected to find 0 rows in system.locations. Found  %d instead", count)
+	}
+}

pkg/util/randutil/rand.go

@@ -30,11 +30,13 @@ func NewPseudoSeed() int64 {
 	return seed
 }
 
-// NewPseudoRand returns an instance of math/rand.Rand seeded from crypto/rand
-// and its seed so we can easily and cheaply generate unique streams of
-// numbers. The created object is not safe for concurrent access.
+// NewPseudoRand returns an instance of math/rand.Rand seeded from the
+// environment variable COCKROACH_RANDOM_SEED.  If that variable is not set,
+// crypto/rand is used to generate a seed. The seed is also returned so we can
+// easily and cheaply generate unique streams of numbers. The created object is
+// not safe for concurrent access.
 func NewPseudoRand() (*rand.Rand, int64) {
-	seed := NewPseudoSeed()
+	seed := envutil.EnvOrDefaultInt64("COCKROACH_RANDOM_SEED", NewPseudoSeed())
 	return rand.New(rand.NewSource(seed)), seed
 }