sql: benchmark the pgbench tcp-b batch

standalone pgbench still provides more detailed instrumentation and
more configurable options, but a pure-go benchmark is useful too,
since its easier to run and integrate with out existing benchmarks.

```
name                              time/op
PgbenchQuery_Cockroach-8          3.16ms ± 1%
PgbenchQuery_Postgres-8            407µs ± 8%
ParallelPgbenchQuery_Cockroach-8  2.45ms ± 9%
ParallelPgbenchQuery_Postgres-8    188µs ±17%

name                              alloc/op
PgbenchQuery_Cockroach-8           230kB ± 0%
ParallelPgbenchQuery_Cockroach-8   246kB ±11%

name                              allocs/op
PgbenchQuery_Cockroach-8           3.48k ± 0%
ParallelPgbenchQuery_Cockroach-8   3.90k ± 7%
```

sql/bench_test.go

@@ -38,6 +38,7 @@ func benchmarkCockroach(b *testing.B, f func(b *testing.B, db *sql.DB)) {
 	defer s.Stop()
 
 	pgUrl, cleanupFn := sqlutils.PGUrl(b, s, security.RootUser, "benchmarkCockroach")
+	pgUrl.Path = "bench"
 	defer cleanupFn()
 
 	db, err := sql.Open("postgres", pgUrl.String())

sql/pgbench/cmd/pgbenchsetup/main.go

@@ -37,7 +37,7 @@ var usage = func() {
 func main() {
 	accounts := flag.Int("accounts", 100000, "number of accounts to create")
 
-	createDb := flag.Bool("createdb", false, "attempt to create named db, dropping first if exists")
+	createDb := flag.Bool("createdb", false, "attempt to create named db, dropping first if exists (must be able to connect to default db to do so).")
 
 	flag.Parse()
 	flag.Usage = usage
@@ -46,41 +46,31 @@ func main() {
 		os.Exit(2)
 	}
 
+	var db *sql.DB
+	var err error
+
 	if *createDb {
 		name := ""
-		parsed, err := url.Parse(flag.Arg(0))
-		if err != nil {
-			panic(err)
+		parsed, parseErr := url.Parse(flag.Arg(0))
+		if parseErr != nil {
+			panic(parseErr)
 		} else if len(parsed.Path) < 2 { // first char is '/'
 			panic("URL must include db name")
 		} else {
 			name = parsed.Path[1:]
 		}
 
-		// Create throw-away connection to create the DB.
-		db, err := sql.Open("postgres", flag.Arg(0))
-		if err != nil {
-			panic(err)
-		}
-		defer db.Close()
-
-		if _, err := db.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", name)); err != nil {
-			panic(fmt.Sprintf("Could not drop db %s: %s\n", name, err.Error()))
-		}
-		if _, err := db.Exec(fmt.Sprintf("CREATE DATABASE %s", name)); err != nil {
-			panic(fmt.Sprintf("Could not create db %s: %s\n", name, err.Error()))
-		} else {
-			fmt.Printf("Created database %s\n", name)
-		}
+		db, err = pgbench.CreateAndConnect(*parsed, name)
+	} else {
+		db, err = sql.Open("postgres", flag.Arg(0))
 	}
-
-	db, err := sql.Open("postgres", flag.Arg(0))
 	if err != nil {
 		panic(err)
 	}
+
 	defer db.Close()
 
-	if err := pgbench.SetupBenchDB(db, *accounts); err != nil {
+	if err := pgbench.SetupBenchDB(db, *accounts, false); err != nil {
 		panic(err)
 	}
 }

sql/pgbench/query.go

@@ -0,0 +1,46 @@
+// Copyright 2016 The Cockroach Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+// implied. See the License for the specific language governing
+// permissions and limitations under the License.
+//
+// Author: David Taylor (david@cockroachlabs.com)
+
+package pgbench
+
+import (
+	"database/sql"
+	"fmt"
+	"math/rand"
+)
+
+// This is the TPC-B(ish) query that pgbench runs.
+// We don't use placeholders because pgwire protocol does not
+// allow multiple statements in prepared queries.
+const tpcbQuery = `BEGIN;
+UPDATE pgbench_accounts SET abalance = abalance + %[1]d WHERE aid = %[2]d;
+SELECT abalance FROM pgbench_accounts WHERE aid = %[2]d;
+UPDATE pgbench_tellers SET tbalance = tbalance + %[1]d WHERE tid = %[3]d;
+UPDATE pgbench_branches SET bbalance = bbalance + %[1]d WHERE bid = %[4]d;
+INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (%[3]d, %[4]d, %[2]d, %[1]d, CURRENT_TIMESTAMP);
+END;` // vars: 1 delta, 2 aid, 3 tid, 4 bid
+
+// RunOne executes one iteration of the query batch that `pgbench` executes.
+func RunOne(db *sql.DB, r *rand.Rand, accounts int) error {
+	account := r.Intn(accounts)
+	delta := r.Intn(5000)
+	teller := r.Intn(tellers)
+	branch := 1
+
+	q := fmt.Sprintf(tpcbQuery, delta, account, teller, branch)
+	_, err := db.Exec(q)
+	return err
+}

sql/pgbench/setup.go

@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"database/sql"
 	"fmt"
+	"net/url"
 )
 
 const schema = `
@@ -58,26 +59,56 @@ CREATE TABLE pgbench_history (
 );
 `
 
+// CreateAndConnect connects and creates the requested DB (dropping
+// if exists) then returns a new connection to the created DB.
+func CreateAndConnect(pgURL url.URL, name string) (*sql.DB, error) {
+	{
+		db, err := sql.Open("postgres", pgURL.String())
+		if err != nil {
+			return nil, err
+		}
+		defer db.Close()
+
+		if _, err := db.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", name)); err != nil {
+			return nil, err
+		}
+
+		if _, err := db.Exec(fmt.Sprintf(`CREATE DATABASE %s`, name)); err != nil {
+			return nil, err
+		}
+	}
+
+	pgURL.Path = name
+
+	db, err := sql.Open("postgres", pgURL.String())
+	if err != nil {
+		return nil, err
+	}
+	return db, nil
+}
+
 // SetupBenchDB sets up a db with the schema and initial data used by `pgbench`.
 // The `-i` flag to `pgbench` is usually used to do this when testing postgres
 // but the statements it generates use postgres-specific flags that cockroach does
 // not support. The queries this script runs are based on a dump of a db created
 // by `pgbench -i`, but sticking to the compatible subset that both cockroach and
 // postgres support.
-func SetupBenchDB(db *sql.DB, accounts int) error {
+func SetupBenchDB(db *sql.DB, accounts int, quiet bool) error {
 	if _, err := db.Exec(schema); err != nil {
 		return err
 	}
-	return populateDB(db, accounts)
+	return populateDB(db, accounts, quiet)
 }
 
-func populateDB(db *sql.DB, accounts int) error {
+const tellers = 10
+
+func populateDB(db *sql.DB, accounts int, quiet bool) error {
 	branches := `INSERT INTO pgbench_branches (bid, bbalance, filler) VALUES (1, 7354, NULL)`
 	if r, err := db.Exec(branches); err != nil {
 		return err
 	} else if x, err := r.RowsAffected(); err != nil {
 		return err
-	} else {
+	} else if !quiet {
 		fmt.Printf("Inserted %d branch records\n", x)
 	}
 
@@ -97,7 +128,7 @@ func populateDB(db *sql.DB, accounts int) error {
 		return err
 	} else if x, err := r.RowsAffected(); err != nil {
 		return err
-	} else {
+	} else if !quiet {
 		fmt.Printf("Inserted %d teller records\n", x)
 	}
 
@@ -124,7 +155,7 @@ func populateDB(db *sql.DB, accounts int) error {
 			return err
 		} else if x, err := r.RowsAffected(); err != nil {
 			return err
-		} else {
+		} else if !quiet {
 			fmt.Printf("Inserted %d account records\n", x)
 		}
 		done += batch
@@ -147,7 +178,7 @@ INSERT INTO pgbench_history VALUES
 		return err
 	} else if x, err := r.RowsAffected(); err != nil {
 		return err
-	} else {
+	} else if !quiet {
 		fmt.Printf("Inserted %d history records\n", x)
 	}
 

sql/pgbench_test.go

@@ -0,0 +1,77 @@
+// Copyright 2016 The Cockroach Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+// implied. See the License for the specific language governing
+// permissions and limitations under the License.
+//
+// Author: David Taylor (david@cockroachlabs.com)
+
+package sql_test
+
+import (
+	"database/sql"
+	"math/rand"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/sql/pgbench"
+)
+
+// Tests a batch of queries very similar to those that that PGBench runs
+// in its TPC-B(ish) mode.
+func runPgbenchQuery(b *testing.B, db *sql.DB) {
+	if err := pgbench.SetupBenchDB(db, 20000, true /*quiet*/); err != nil {
+		b.Fatal(err)
+	}
+	src := rand.New(rand.NewSource(5432))
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if err := pgbench.RunOne(db, src, 20000); err != nil {
+			b.Fatal(err)
+		}
+	}
+	b.StopTimer()
+}
+
+// Tests a batch of queries very similar to those that that PGBench runs
+// in its TPC-B(ish) mode.
+func runPgbenchQueryParallel(b *testing.B, db *sql.DB) {
+	if err := pgbench.SetupBenchDB(db, 20000, true /*quiet*/); err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		src := rand.New(rand.NewSource(5432))
+		for pb.Next() {
+			if err := pgbench.RunOne(db, src, 20000); err != nil {
+				// TODO(dt): handle retry/aborted correctly
+				// b.Fatal(err)
+			}
+		}
+	})
+	b.StopTimer()
+}
+
+func BenchmarkPgbenchQuery_Cockroach(b *testing.B) {
+	benchmarkCockroach(b, runPgbenchQuery)
+}
+
+func BenchmarkPgbenchQuery_Postgres(b *testing.B) {
+	benchmarkPostgres(b, runPgbenchQuery)
+}
+
+func BenchmarkParallelPgbenchQuery_Cockroach(b *testing.B) {
+	benchmarkCockroach(b, runPgbenchQueryParallel)
+}
+
+func BenchmarkParallelPgbenchQuery_Postgres(b *testing.B) {
+	benchmarkPostgres(b, runPgbenchQueryParallel)
+}