Real world cockroach benchmark

main.go

package main

import (
	"context"
	"os"
	"strconv"
	"sync"
	"time"

	otgorm "github.com/echo-health/opentracing-gorm"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/postgres"
	_ "github.com/mattes/migrate/source/file"
	"github.com/prometheus/common/log"
	uuid "github.com/satori/go.uuid"
	"github.com/segmentio/ksuid"
	"github.com/wawandco/fako"
)

var (
	defaultConnString = "host=localhost user=postgres dbname=postgres sslmode=disable"
	defaultMaxConn    = 20
	defaultPoolName   = "default"
	lock              = &sync.Mutex{}
	conns             = map[string]*gorm.DB{}
)

const schema = `
CREATE TABLE IF NOT EXISTS users (
	id STRING NOT NULL,
	firebase_id STRING NOT NULL,
	practice_id STRING NOT NULL,
	title STRING NOT NULL,
	first_name STRING NOT NULL,
	last_name STRING NOT NULL,
	gender STRING NOT NULL,
	date_of_birth TIMESTAMP NOT NULL,
	create_time TIMESTAMP NULL DEFAULT now(),
	update_time TIMESTAMP NULL DEFAULT now(),
	CONSTRAINT "primary" PRIMARY KEY (id ASC),
	UNIQUE INDEX idx_users_firebase_id (firebase_id ASC)
);

CREATE TABLE IF NOT EXISTS registered_addresses (
	user_id STRING NOT NULL,
	id STRING NOT NULL,
	line1 STRING NOT NULL,
	line2 STRING NULL,
	city STRING NOT NULL,
	postcode STRING NOT NULL,
	country STRING NOT NULL,
	lat DECIMAL NOT NULL,
	lng DECIMAL NOT NULL,
	CONSTRAINT "primary" PRIMARY KEY (user_id, id),
	CONSTRAINT fk_user FOREIGN KEY (user_id) REFERENCES users
) INTERLEAVE IN PARENT users (user_id);

CREATE TABLE IF NOT EXISTS user_relationships (
	user_id STRING NOT NULL,
	account_id STRING NOT NULL,
	relationship STRING NOT NULL,
	create_time TIMESTAMP NULL DEFAULT now(),
	update_time TIMESTAMP NULL DEFAULT now(),
	CONSTRAINT "primary" PRIMARY KEY (user_id, account_id),
	INDEX idx_users_accounts_id (account_id ASC),
	CONSTRAINT fk_user FOREIGN KEY (user_id) REFERENCES users
) INTERLEAVE IN PARENT users (user_id);

CREATE TABLE IF NOT EXISTS user_exemptions (
	id STRING NOT NULL,
	user_id STRING NOT NULL,
	code STRING NOT NULL,
	status STRING NOT NULL,
	url STRING NULL,
	expiry TIMESTAMP NULL DEFAULT now(),
	create_time TIMESTAMP NULL DEFAULT now(),
	update_time TIMESTAMP NULL DEFAULT now(),
	CONSTRAINT "primary" PRIMARY KEY (user_id, id),
	CONSTRAINT fk_user FOREIGN KEY (user_id) REFERENCES users
) INTERLEAVE IN PARENT users (user_id);
`

type User struct {
	ID                string
	PracticeID        string
	FirebaseID        string
	Title             string `fako:"title"`
	FirstName         string `fako:"first_name"`
	LastName          string `fako:"last_name"`
	Gender            string
	DateOfBirth       time.Time
	CreatedAt         time.Time `gorm:"column:create_time"`
	UpdatedAt         time.Time `gorm:"column:update_time"`
	Relationships     []*UserRelationship
	Exemptions        []*UserExemption
	RegisteredAddress RegisteredAddress
}

type UserRelationship struct {
	UserID       string    `gorm:"primary_key"`
	AccountID    string    `gorm:"primary_key" fako:"simple_password"`
	Relationship string    `fako:"first_name"`
	CreatedAt    time.Time `gorm:"column:create_time"`
	UpdatedAt    time.Time `gorm:"column:update_time"`
}

type RegisteredAddress struct {
	ID       string
	UserID   string
	Line1    string `fako:"street_address"`
	Line2    string `fako:"street"`
	City     string `fako:"city"`
	Postcode string `fako:"zip"`
	Country  string `fako:"country"`
	Lat      float64
	Lng      float64
}

type UserExemption struct {
	ID        string
	UserID    string
	URL       string `fako:"characters"`
	Code      string
	Status    string
	Expiry    *time.Time
	CreatedAt time.Time `gorm:"column:create_time"`
	UpdatedAt time.Time `gorm:"column:update_time"`
}

func UserID() string {
	if os.Getenv("USE_UUID") != "" {
		u, _ := uuid.NewV4()
		return u.String()
	}

	return "user_" + ksuid.New().String()
}

func ExemptionID() string {
	if os.Getenv("USE_UUID") != "" {
		u, _ := uuid.NewV4()
		return u.String()
	}

	return "user_ex_" + ksuid.New().String()
}

func NominationRequestID() string {
	if os.Getenv("USE_UUID") != "" {
		u, _ := uuid.NewV4()
		return u.String()
	}

	return "user_nom_" + ksuid.New().String()
}

func AddrID() string {
	if os.Getenv("USE_UUID") != "" {
		u, _ := uuid.NewV4()
		return u.String()
	}

	return "addr_" + ksuid.New().String()
}

func (p *User) BeforeCreate(scope *gorm.Scope) error {
	if p.ID == "" {
		p.ID = UserID()
		scope.SetColumn("ID", p.ID)
	}

	if p.FirebaseID == "" {
		p.FirebaseID = p.ID
	}

	return nil
}

func (p *RegisteredAddress) BeforeCreate(scope *gorm.Scope) error {
	if p.ID == "" {
		p.ID = AddrID()
		return scope.SetColumn("ID", p.ID)
	}

	return nil
}

func (p *UserExemption) BeforeCreate(scope *gorm.Scope) error {
	if p.ID == "" {
		p.ID = ExemptionID()
		return scope.SetColumn("ID", p.ID)
	}

	return nil
}

func GetDB(ctx context.Context) (*gorm.DB, error) {
	return GetDBFromPool(ctx, defaultPoolName)
}

func GetDBFromPool(ctx context.Context, pool string) (*gorm.DB, error) {
	lock.Lock()
	defer lock.Unlock()

	if conns[pool] != nil {
		db := otgorm.SetSpanToGorm(ctx, conns[pool])
		return db, nil
	}

	connString := defaultConnString
	if os.Getenv("POSTGRES_URL") != "" {
		connString = os.Getenv("POSTGRES_URL")
	}

	log.Debugf("Connecting to PG at %s for pool %s", connString, pool)
	var err error
	conn, err := connect(connString)
	if err != nil {
		db := otgorm.SetSpanToGorm(ctx, conn)
		return db, err
	}

	conns[pool] = conn
	return conns[pool], err
}

func connect(connString string) (*gorm.DB, error) {
	driver := "postgres"
	if os.Getenv("CLOUDSQL") == "true" {
		driver = "cloudsqlpostgres"
	}

	conn, err := gorm.Open(driver, connString)

	if err != nil {
		return nil, err
	}

	maxConns := defaultMaxConn
	if os.Getenv("DB_MAX_CONNS") != "" {
		i, err := strconv.Atoi(os.Getenv("DB_MAX_CONNS"))
		if err != nil {
			return nil, err
		}

		maxConns = i
	}

	conn.DB().SetMaxOpenConns(maxConns)

	if os.Getenv("DEBUG_SQL") == "true" {
		conn.LogMode(true)
	}

	otgorm.AddGormCallbacks(conn)

	conn = conn.Set("gorm:auto_preload", true)
	return conn, nil
}

func Get(ctx context.Context, id string) (*User, error) {
	db, err := GetDB(ctx)
	if err != nil {
		return nil, err
	}

	var p User
	err = db.Where("id = ? OR firebase_id = ?", id, id).
		First(&p).Error
	return &p, err
}

func Create(ctx context.Context, new *User) error {
	db, err := GetDBFromPool(ctx, "create")
	if err != nil {
		return err
	}

	old, err := Get(ctx, new.FirebaseID)
	if err == gorm.ErrRecordNotFound {
		return db.Create(new).Error
	}

	if err != nil {
		return err
	}

	tx := db.Begin()

	// first delete any relationships that have been removed from firebase
	for _, rel := range old.Relationships {
		inFirebase := false
		for _, a := range new.Relationships {
			if a.AccountID == rel.AccountID {
				inFirebase = true
			}
		}

		if !inFirebase {
			if err = tx.Delete(rel).Error; err != nil {
				tx.Rollback()
				return err
			}
		}
	}

	new.ID = old.ID
	new.RegisteredAddress.ID = old.RegisteredAddress.ID
	err = tx.Save(&new).Error
	if err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit().Error
}

func worker(id int, jobs <-chan *User) {
	log.Infof("Booting worker %d", id)
	for j := range jobs {
		err := Create(context.Background(), j)
		if err != nil {
			log.Fatal(err)
		}
	}
}

func main() {
	db, err := GetDB(context.Background())
	if err != nil {
		log.Fatal(err)
	}

	err = db.Exec(schema).Error
	if err != nil {
		log.Fatal(err)
	}

	jobs := make(chan *User, 100)

	for w := 1; w <= 20; w++ {
		go worker(w, jobs)
	}

	for i := 0; i < 10000; i++ {
		u := &User{}
		fako.Fill(u)

		ra := &RegisteredAddress{}
		fako.Fill(ra)
		u.RegisteredAddress = *ra

		for i := 0; i < 2; i++ {
			r := &UserRelationship{}
			fako.Fill(r)
			u.Relationships = append(u.Relationships, r)
		}

		for i := 0; i < 2; i++ {
			r := &UserExemption{}
			fako.Fill(r)
			u.Exemptions = append(u.Exemptions, r)
		}

		jobs <- u
	}
}