movr.go

// Copyright 2019 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package movr

import (
	"context"
	gosql "database/sql"
	"math"
	"strings"
	"sync"
	"time"

	"github.com/cockroachdb/cockroach/pkg/util/uuid"
	"github.com/cockroachdb/cockroach/pkg/workload"
	"github.com/cockroachdb/cockroach/pkg/workload/faker"
	"github.com/cockroachdb/cockroach/pkg/workload/histogram"
	"github.com/pkg/errors"
	"github.com/spf13/pflag"
	"golang.org/x/exp/rand"
)

// Indexes into the slice returned by `Tables`.
const (
	TablesUsersIdx                    = 0
	TablesVehiclesIdx                 = 1
	TablesRidesIdx                    = 2
	TablesVehicleLocationHistoriesIdx = 3
	TablesPromoCodesIdx               = 4
	TablesUserPromoCodesIdx           = 5
)

const movrUsersSchema = `(
  id UUID NOT NULL,
  city VARCHAR NOT NULL,
  name VARCHAR NULL,
  address VARCHAR NULL,
  credit_card VARCHAR NULL,
  PRIMARY KEY (city ASC, id ASC)
)`
const movrVehiclesSchema = `(
  id UUID NOT NULL,
  city VARCHAR NOT NULL,
  type VARCHAR NULL,
  owner_id UUID NULL,
  creation_time TIMESTAMP NULL,
  status VARCHAR NULL,
  current_location VARCHAR NULL,
  ext JSONB NULL,
  PRIMARY KEY (city ASC, id ASC),
  INDEX vehicles_auto_index_fk_city_ref_users (city ASC, owner_id ASC)
)`
const movrRidesSchema = `(
  id UUID NOT NULL,
  city VARCHAR NOT NULL,
  vehicle_city VARCHAR NULL,
  rider_id UUID NULL,
  vehicle_id UUID NULL,
  start_address VARCHAR NULL,
  end_address VARCHAR NULL,
  start_time TIMESTAMP NULL,
  end_time TIMESTAMP NULL,
  revenue DECIMAL(10,2) NULL,
  PRIMARY KEY (city ASC, id ASC),
  INDEX rides_auto_index_fk_city_ref_users (city ASC, rider_id ASC),
  INDEX rides_auto_index_fk_vehicle_city_ref_vehicles (vehicle_city ASC, vehicle_id ASC),
  CONSTRAINT check_vehicle_city_city CHECK (vehicle_city = city)
)`
const movrVehicleLocationHistoriesSchema = `(
  city VARCHAR NOT NULL,
  ride_id UUID NOT NULL,
  "timestamp" TIMESTAMP NOT NULL,
  lat FLOAT8 NULL,
  long FLOAT8 NULL,
  PRIMARY KEY (city ASC, ride_id ASC, "timestamp" ASC)
)`
const movrPromoCodesSchema = `(
  code VARCHAR NOT NULL,
  description VARCHAR NULL,
  creation_time TIMESTAMP NULL,
  expiration_time TIMESTAMP NULL,
  rules JSONB NULL,
  PRIMARY KEY (code ASC)
)`
const movrUserPromoCodesSchema = `(
  city VARCHAR NOT NULL,
  user_id UUID NOT NULL,
  code VARCHAR NOT NULL,
  "timestamp" TIMESTAMP NULL,
  usage_count INT NULL,
  PRIMARY KEY (city ASC, user_id ASC, code ASC)
)`

const (
	timestampFormat = "2006-01-02 15:04:05.999999-07:00"
)

var cities = []struct {
	city     string
	locality string
}{
	{city: "new york", locality: "us_east"},
	{city: "boston", locality: "us_east"},
	{city: "washington dc", locality: "us_east"},
	{city: "seattle", locality: "us_west"},
	{city: "san francisco", locality: "us_west"},
	{city: "los angeles", locality: "us_west"},
	// The demo setup we have is a 9 node, 3 region cluster.
	//{city: "chicago", locality: "us_central"},
	//{city: "detroit", locality: "us_central"},
	//{city: "minneapolis", locality: "us_central"},
	{city: "amsterdam", locality: "eu_west"},
	{city: "paris", locality: "eu_west"},
	{city: "rome", locality: "eu_west"},
}

type movr struct {
	flags     workload.Flags
	connFlags *workload.ConnFlags

	seed                              uint64
	users, vehicles, rides, histories cityDistributor
	numPromoCodes                     int

	creationTime time.Time

	fakerOnce sync.Once
	faker     faker.Faker
}

func init() {
	workload.Register(movrMeta)
}

var movrMeta = workload.Meta{
	Name:         `movr`,
	Description:  `MovR is a fictional ride sharing company`,
	Version:      `1.0.0`,
	PublicFacing: true,
	New: func() workload.Generator {
		g := &movr{}
		g.flags.FlagSet = pflag.NewFlagSet(`movr`, pflag.ContinueOnError)
		g.flags.Uint64Var(&g.seed, `seed`, 1, `Key hash seed.`)
		g.flags.IntVar(&g.users.numRows, `num-users`, 50, `Initial number of users.`)
		g.flags.IntVar(&g.vehicles.numRows, `num-vehicles`, 15, `Initial number of vehicles.`)
		g.flags.IntVar(&g.rides.numRows, `num-rides`, 500, `Initial number of rides.`)
		g.flags.IntVar(&g.histories.numRows, `num-histories`, 1000,
			`Initial number of ride location histories.`)
		g.flags.IntVar(&g.numPromoCodes, `num-promo-codes`, 1000, `Initial number of promo codes.`)
		g.connFlags = workload.NewConnFlags(&g.flags)
		g.creationTime = time.Date(2019, 1, 2, 3, 4, 5, 6, time.UTC)
		return g
	},
}

// Meta implements the Generator interface.
func (*movr) Meta() workload.Meta { return movrMeta }

// Flags implements the Flagser interface.
func (g *movr) Flags() workload.Flags { return g.flags }

// Hooks implements the Hookser interface.
func (g *movr) Hooks() workload.Hooks {
	return workload.Hooks{
		Validate: func() error {
			// Force there to be at least one user/vehicle/ride/history per city.
			// Otherwise, some cities will be empty, which means we can't construct
			// the FKs we need.
			if g.users.numRows < len(cities) {
				return errors.Errorf(`at least %d users are required`, len(cities))
			}
			if g.vehicles.numRows < len(cities) {
				return errors.Errorf(`at least %d vehicles are required`, len(cities))
			}
			if g.rides.numRows < len(cities) {
				return errors.Errorf(`at least %d rides are required`, len(cities))
			}
			if g.histories.numRows < len(cities) {
				return errors.Errorf(`at least %d histories are required`, len(cities))
			}
			return nil
		},
		PostLoad: func(db *gosql.DB) error {
			fkStmts := []string{
				`ALTER TABLE vehicles ADD FOREIGN KEY ` +
					`(city, owner_id) REFERENCES users (city, id)`,
				`ALTER TABLE rides ADD FOREIGN KEY ` +
					`(city, rider_id) REFERENCES users (city, id)`,
				`ALTER TABLE rides ADD FOREIGN KEY ` +
					`(vehicle_city, vehicle_id) REFERENCES vehicles (city, id)`,
				`ALTER TABLE vehicle_location_histories ADD FOREIGN KEY ` +
					`(city, ride_id) REFERENCES rides (city, id)`,
				`ALTER TABLE user_promo_codes ADD FOREIGN KEY ` +
					`(city, user_id) REFERENCES users (city, id)`,
			}

			for _, fkStmt := range fkStmts {
				if _, err := db.Exec(fkStmt); err != nil {
					// If the statement failed because the fk already exists,
					// ignore it. Return the error for any other reason.
					const duplicateFKErr = "columns cannot be used by multiple foreign key constraints"
					if !strings.Contains(err.Error(), duplicateFKErr) {
						return err
					}
				}
			}

			// TODO(dan): Partitions.
			return nil
		},
	}
}

// Tables implements the Generator interface.
func (g *movr) Tables() []workload.Table {
	g.fakerOnce.Do(func() {
		g.faker = faker.NewFaker()
	})
	tables := make([]workload.Table, 6)
	tables[TablesUsersIdx] = workload.Table{
		Name:   `users`,
		Schema: movrUsersSchema,
		InitialRows: workload.Tuples(
			g.users.numRows,
			g.movrUsersInitialRow,
		),
	}
	tables[TablesVehiclesIdx] = workload.Table{
		Name:   `vehicles`,
		Schema: movrVehiclesSchema,
		InitialRows: workload.Tuples(
			g.vehicles.numRows,
			g.movrVehiclesInitialRow,
		),
	}
	tables[TablesRidesIdx] = workload.Table{
		Name:   `rides`,
		Schema: movrRidesSchema,
		InitialRows: workload.Tuples(
			g.rides.numRows,
			g.movrRidesInitialRow,
		),
	}
	tables[TablesVehicleLocationHistoriesIdx] = workload.Table{
		Name:   `vehicle_location_histories`,
		Schema: movrVehicleLocationHistoriesSchema,
		InitialRows: workload.Tuples(
			g.histories.numRows,
			g.movrVehicleLocationHistoriesInitialRow,
		),
	}
	tables[TablesPromoCodesIdx] = workload.Table{
		Name:   `promo_codes`,
		Schema: movrPromoCodesSchema,
		InitialRows: workload.Tuples(
			g.numPromoCodes,
			g.movrPromoCodesInitialRow,
		),
	}
	tables[TablesUserPromoCodesIdx] = workload.Table{
		Name:   `user_promo_codes`,
		Schema: movrUserPromoCodesSchema,
		InitialRows: workload.Tuples(
			0,
			func(_ int) []interface{} { panic(`unimplemented`) },
		),
	}
	return tables
}

// cityDistributor deterministically maps each of numRows to a city. It also
// maps a city back to a range of rows. This allows the generator functions
// below to select random rows from the same city in another table. numRows is
// required to be at least `len(cities)`.
type cityDistributor struct {
	numRows int
}

func (d cityDistributor) cityForRow(rowIdx int) int {
	if d.numRows < len(cities) {
		panic(errors.Errorf(`a minimum of %d rows are required got %d`, len(cities), d.numRows))
	}
	numPerCity := float64(d.numRows) / float64(len(cities))
	cityIdx := int(float64(rowIdx) / numPerCity)
	return cityIdx
}

func (d cityDistributor) rowsForCity(cityIdx int) (min, max int) {
	if d.numRows < len(cities) {
		panic(errors.Errorf(`a minimum of %d rows are required got %d`, len(cities), d.numRows))
	}
	numPerCity := float64(d.numRows) / float64(len(cities))
	min = int(math.Ceil(float64(cityIdx) * numPerCity))
	max = int(math.Ceil(float64(cityIdx+1) * numPerCity))
	if min >= d.numRows {
		min = d.numRows
	}
	if max >= d.numRows {
		max = d.numRows
	}
	return min, max
}

func (d cityDistributor) randRowInCity(rng *rand.Rand, cityIdx int) int {
	min, max := d.rowsForCity(cityIdx)
	return min + rng.Intn(max-min)
}

func (g *movr) movrUsersInitialRow(rowIdx int) []interface{} {
	rng := rand.New(rand.NewSource(g.seed + uint64(rowIdx)))
	cityIdx := g.users.cityForRow(rowIdx)
	city := cities[cityIdx]

	// Make evenly-spaced UUIDs sorted in the same order as the rows.
	var id uuid.UUID
	id.DeterministicV4(uint64(rowIdx), uint64(g.users.numRows))

	return []interface{}{
		id.String(),                // id
		city.city,                  // city
		g.faker.Name(rng),          // name
		g.faker.StreetAddress(rng), // address
		randCreditCard(rng),        // credit_card
	}
}

func (g *movr) movrVehiclesInitialRow(rowIdx int) []interface{} {
	rng := rand.New(rand.NewSource(g.seed + uint64(rowIdx)))
	cityIdx := g.vehicles.cityForRow(rowIdx)
	city := cities[cityIdx]

	// Make evenly-spaced UUIDs sorted in the same order as the rows.
	var id uuid.UUID
	id.DeterministicV4(uint64(rowIdx), uint64(g.vehicles.numRows))

	vehicleType := randVehicleType(rng)
	ownerRowIdx := g.users.randRowInCity(rng, cityIdx)
	ownerID := g.movrUsersInitialRow(ownerRowIdx)[0]

	return []interface{}{
		id.String(),                            // id
		city.city,                              // city
		vehicleType,                            // type
		ownerID,                                // owner_id
		g.creationTime.Format(timestampFormat), // creation_time
		randVehicleStatus(rng),                 // status
		g.faker.StreetAddress(rng),             // current_location
		randVehicleMetadata(rng, vehicleType),  // ext
	}
}

func (g *movr) movrRidesInitialRow(rowIdx int) []interface{} {
	rng := rand.New(rand.NewSource(g.seed + uint64(rowIdx)))
	cityIdx := g.rides.cityForRow(rowIdx)
	city := cities[cityIdx]

	// Make evenly-spaced UUIDs sorted in the same order as the rows.
	var id uuid.UUID
	id.DeterministicV4(uint64(rowIdx), uint64(g.rides.numRows))

	riderRowIdx := g.users.randRowInCity(rng, cityIdx)
	riderID := g.movrUsersInitialRow(riderRowIdx)[0]
	vehicleRowIdx := g.vehicles.randRowInCity(rng, cityIdx)
	vehicleID := g.movrVehiclesInitialRow(vehicleRowIdx)[0]
	startTime := g.creationTime.Add(-time.Duration(rng.Intn(30)) * 24 * time.Hour)
	endTime := startTime.Add(time.Duration(rng.Intn(60)) * time.Hour)

	return []interface{}{
		id.String(),                       // id
		city.city,                         // city
		city.city,                         // vehicle_city
		riderID,                           // rider_id
		vehicleID,                         // vehicle_id
		g.faker.StreetAddress(rng),        // start_address
		g.faker.StreetAddress(rng),        // end_address
		startTime.Format(timestampFormat), // start_time
		endTime.Format(timestampFormat),   // end_time
		rng.Intn(100),                     // revenue
	}
}

func (g *movr) movrVehicleLocationHistoriesInitialRow(rowIdx int) []interface{} {
	rng := rand.New(rand.NewSource(g.seed + uint64(rowIdx)))
	cityIdx := g.histories.cityForRow(rowIdx)
	city := cities[cityIdx]

	rideRowIdx := g.rides.randRowInCity(rng, cityIdx)
	rideID := g.movrRidesInitialRow(rideRowIdx)[0]
	time := g.creationTime.Add(time.Duration(rowIdx) * time.Millisecond)
	lat, long := randLatLong(rng)

	return []interface{}{
		city.city,                    // city
		rideID,                       // ride_id,
		time.Format(timestampFormat), // timestamp
		lat,                          // lat
		long,                         // long
	}
}

func (g *movr) movrPromoCodesInitialRow(rowIdx int) []interface{} {
	rng := rand.New(rand.NewSource(g.seed + uint64(rowIdx)))
	code := strings.ToLower(strings.Join(g.faker.Words(rng, 3), `_`))
	description := g.faker.Paragraph(rng)
	expirationTime := g.creationTime.Add(time.Duration(rng.Intn(30)) * 24 * time.Hour)
	// TODO(dan): This is nil in the reference impl, is that intentional?
	creationTime := expirationTime.Add(-time.Duration(rng.Intn(30)) * 24 * time.Hour)
	const rulesJSON = `{"type": "percent_discount", "value": "10%"}`

	return []interface{}{
		code,           // code
		description,    // description
		creationTime,   // creation_time
		expirationTime, // expiration_time
		rulesJSON,      // rules
	}
}

type rideInfo struct {
	id   string
	city string
}

// Ops implements the Opser interface
func (g *movr) Ops(urls []string, reg *histogram.Registry) (workload.QueryLoad, error) {
	sqlDatabase, err := workload.SanitizeUrls(g, g.connFlags.DBOverride, urls)
	if err != nil {
		return workload.QueryLoad{}, err
	}
	db, err := gosql.Open(`postgres`, strings.Join(urls, ` `))
	if err != nil {
		return workload.QueryLoad{}, err
	}

	ql := workload.QueryLoad{SQLDatabase: sqlDatabase}

	rng := rand.New(rand.NewSource(g.seed))
	readPercentage := 0.90
	activeRides := []rideInfo{}

	getRandomUser := func(city string) (string, error) {
		id, err := uuid.NewV4()
		if err != nil {
			return "", err
		}
		var user string
		q := `
		SELECT
			IFNULL(a, b)
		FROM
			(
				SELECT
					(SELECT id FROM users WHERE city = $1 AND id > $2 ORDER BY id LIMIT 1)
						AS a,
					(SELECT id FROM users WHERE city = $1 ORDER BY id LIMIT 1) AS b
			);
		`
		err = db.QueryRow(q, city, id.String()).Scan(&user)
		return user, err
	}

	getRandomPromoCode := func() (string, error) {
		id, err := uuid.NewV4()
		if err != nil {
			return "", err
		}
		q := `
		SELECT
			IFNULL(a, b)
		FROM
			(
				SELECT
					(SELECT code FROM promo_codes WHERE code > $1 ORDER BY code LIMIT 1)
						AS a,
					(SELECT code FROM promo_codes ORDER BY code LIMIT 1) AS b
			);
		`
		var code string
		err = db.QueryRow(q, id.String()).Scan(&code)
		return code, err
	}

	getRandomVehicle := func(city string) (string, error) {
		id, err := uuid.NewV4()
		if err != nil {
			return "", err
		}
		q := `
		SELECT
			IFNULL(a, b)
		FROM
			(
				SELECT
					(SELECT id FROM vehicles WHERE city = $1 AND id > $2 ORDER BY id LIMIT 1)
						AS a,
					(SELECT id FROM vehicles WHERE city = $1 ORDER BY id LIMIT 1) AS b
			);
		`
		var vehicle string
		err = db.QueryRow(q, city, id.String()).Scan(&vehicle)
		return vehicle, err
	}

	movrQuerySimulation := func(ctx context.Context) error {
		activeCity := randCity(rng)
		if rng.Float64() <= readPercentage {
			q := `SELECT city, id FROM vehicles WHERE city = $1`
			_, err := db.Exec(q, activeCity)
			return err
		}
		// Simulate vehicle location updates.
		for i, ride := range activeRides {
			if i >= 10 {
				break
			}
			lat, long := randLatLong(rng)
			q := `UPSERT INTO vehicle_location_histories VALUES ($1, $2, now(), $3, $4)`
			_, err := db.Exec(q, ride.city, ride.id, lat, long)
			if err != nil {
				return err
			}
		}

		id, err := uuid.NewV4()
		if err != nil {
			return err
		}

		// Do write operations.
		if rng.Float64() < 0.03 {
			q := `INSERT INTO promo_codes VALUES ($1, NULL, NULL, NULL, NULL)`
			_, err = db.Exec(q, id.String())
			return err
		} else if rng.Float64() < 0.1 {
			// Apply a promo code to an account.
			user, err := getRandomUser(activeCity)
			if err != nil {
				return err
			}

			code, err := getRandomPromoCode()
			if err != nil {
				return err
			}

			// See if the promo code has been used.
			var count int
			q := `SELECT count(*) FROM user_promo_codes WHERE city = $1 AND user_id = $2 AND code = $3`
			err = db.QueryRow(q, activeCity, user, code).Scan(&count)
			if err != nil {
				return err
			}

			// If is has not been, apply the promo code.
			if count == 0 {
				q = `INSERT INTO user_promo_codes VALUES ($1, $2, $3, NULL, NULL)`
				_, err = db.Exec(q, activeCity, user, code)
				return err
			}
			return nil
		} else if rng.Float64() < 0.3 {
			q := `INSERT INTO users VALUES ($1, $2, NULL, NULL, NULL)`
			_, err = db.Exec(q, id.String(), activeCity)
			return err
		} else if rng.Float64() < 0.1 {
			// Simulate adding a new vehicle to the population.
			ownerID, err := getRandomUser(activeCity)
			if err != nil {
				return err
			}

			typ := randVehicleType(rng)
			q := `INSERT INTO vehicles VALUES ($1, $2, $3, $4, NULL, NULL, NULL, NULL)`
			_, err = db.Exec(q, id.String(), activeCity, typ, ownerID)
			return err
		} else if rng.Float64() < 0.5 {
			// Simulate a user starting a ride.
			rider, err := getRandomUser(activeCity)
			if err != nil {
				return err
			}

			vehicle, err := getRandomVehicle(activeCity)
			if err != nil {
				return err
			}

			q := `INSERT INTO rides VALUES ($1, $2, $2, $3, $4, $5, NULL, now(), NULL, NULL)`
			_, err = db.Exec(q, id.String(), activeCity, rider, vehicle, g.faker.StreetAddress(rng))
			if err != nil {
				return err
			}
			activeRides = append(activeRides, rideInfo{id.String(), activeCity})
			return err
		} else {
			// Simulate a ride ending.
			if len(activeRides) > 1 {
				ride := activeRides[0]
				activeRides = activeRides[1:]
				q := `UPDATE rides SET end_address = $3, end_time = now() WHERE city = $1 AND id = $2`
				_, err := db.Exec(q, ride.city, ride.id, g.faker.StreetAddress(rng))
				return err
			}
		}

		return nil
	}

	ql.WorkerFns = append(ql.WorkerFns, movrQuerySimulation)

	return ql, nil
}