github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/bench/foreachdb.go (about)

     1  // Copyright 2017 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package bench
    12  
    13  import (
    14  	"context"
    15  	gosql "database/sql"
    16  	"fmt"
    17  	"net"
    18  	"net/url"
    19  	"reflect"
    20  	"runtime"
    21  	"strings"
    22  	"testing"
    23  
    24  	"github.com/cockroachdb/cockroach/pkg/base"
    25  	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
    26  	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
    27  	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
    28  	_ "github.com/go-sql-driver/mysql" // registers the MySQL driver to gosql
    29  	_ "github.com/lib/pq"              // registers the pg driver to gosql
    30  )
    31  
    32  // BenchmarkFn is a function that runs a benchmark using the given SQLRunner.
    33  type BenchmarkFn func(b *testing.B, db *sqlutils.SQLRunner)
    34  
    35  func benchmarkCockroach(b *testing.B, f BenchmarkFn) {
    36  	s, db, _ := serverutils.StartServer(
    37  		b, base.TestServerArgs{UseDatabase: "bench"})
    38  	defer s.Stopper().Stop(context.TODO())
    39  
    40  	if _, err := db.Exec(`CREATE DATABASE bench`); err != nil {
    41  		b.Fatal(err)
    42  	}
    43  
    44  	f(b, sqlutils.MakeSQLRunner(db))
    45  }
    46  
    47  func benchmarkMultinodeCockroach(b *testing.B, f BenchmarkFn) {
    48  	tc := testcluster.StartTestCluster(b, 3,
    49  		base.TestClusterArgs{
    50  			ReplicationMode: base.ReplicationAuto,
    51  			ServerArgs: base.TestServerArgs{
    52  				UseDatabase: "bench",
    53  			},
    54  		})
    55  	if _, err := tc.Conns[0].Exec(`CREATE DATABASE bench`); err != nil {
    56  		b.Fatal(err)
    57  	}
    58  	defer tc.Stopper().Stop(context.TODO())
    59  
    60  	f(b, sqlutils.MakeRoundRobinSQLRunner(tc.Conns[0], tc.Conns[1], tc.Conns[2]))
    61  }
    62  
    63  func benchmarkPostgres(b *testing.B, f BenchmarkFn) {
    64  	// Note: the following uses SSL. To run this, make sure your local
    65  	// Postgres server has SSL enabled. To use Cockroach's checked-in
    66  	// testing certificates for Postgres' SSL, first determine the
    67  	// location of your Postgres server's configuration file:
    68  	// ```
    69  	// $ psql -h localhost -p 5432 -c 'SHOW config_file'
    70  	//                config_file
    71  	// -----------------------------------------
    72  	//  /usr/local/var/postgres/postgresql.conf
    73  	// (1 row)
    74  	//```
    75  	//
    76  	// Now open this file and set the following values:
    77  	// ```
    78  	// $ grep ^ssl /usr/local/var/postgres/postgresql.conf
    79  	// ssl = on # (change requires restart)
    80  	// ssl_cert_file = '$GOPATH/src/github.com/cockroachdb/cockroach/pkg/security/securitytest/test_certs/node.crt' # (change requires restart)
    81  	// ssl_key_file = '$GOPATH/src/github.com/cockroachdb/cockroach/pkg/security/securitytest/test_certs/node.key' # (change requires restart)
    82  	// ssl_ca_file = '$GOPATH/src/github.com/cockroachdb/cockroach/pkg/security/securitytest/test_certs/ca.crt' # (change requires restart)
    83  	// ```
    84  	// Where `$GOPATH/src/github.com/cockroachdb/cockroach`
    85  	// is replaced with your local Cockroach source directory.
    86  	// Be sure to restart Postgres for this to take effect.
    87  
    88  	pgURL := url.URL{
    89  		Scheme:   "postgres",
    90  		Host:     "localhost:5432",
    91  		RawQuery: "sslmode=require&dbname=postgres",
    92  	}
    93  	if conn, err := net.Dial("tcp", pgURL.Host); err != nil {
    94  		b.Skipf("unable to connect to postgres server on %s: %s", pgURL.Host, err)
    95  	} else {
    96  		conn.Close()
    97  	}
    98  
    99  	db, err := gosql.Open("postgres", pgURL.String())
   100  	if err != nil {
   101  		b.Fatal(err)
   102  	}
   103  	defer db.Close()
   104  
   105  	r := sqlutils.MakeSQLRunner(db)
   106  	r.Exec(b, `CREATE SCHEMA IF NOT EXISTS bench`)
   107  
   108  	f(b, r)
   109  }
   110  
   111  func benchmarkMySQL(b *testing.B, f BenchmarkFn) {
   112  	const addr = "localhost:3306"
   113  	if conn, err := net.Dial("tcp", addr); err != nil {
   114  		b.Skipf("unable to connect to mysql server on %s: %s", addr, err)
   115  	} else {
   116  		conn.Close()
   117  	}
   118  
   119  	db, err := gosql.Open("mysql", fmt.Sprintf("root@tcp(%s)/", addr))
   120  	if err != nil {
   121  		b.Fatal(err)
   122  	}
   123  	defer db.Close()
   124  
   125  	r := sqlutils.MakeSQLRunner(db)
   126  	r.Exec(b, `CREATE DATABASE IF NOT EXISTS bench`)
   127  
   128  	f(b, r)
   129  }
   130  
   131  // ForEachDB iterates the given benchmark over multiple database engines.
   132  func ForEachDB(b *testing.B, fn BenchmarkFn) {
   133  	for _, dbFn := range []func(*testing.B, BenchmarkFn){
   134  		benchmarkCockroach,
   135  		benchmarkMultinodeCockroach,
   136  		benchmarkPostgres,
   137  		benchmarkMySQL,
   138  	} {
   139  		dbName := runtime.FuncForPC(reflect.ValueOf(dbFn).Pointer()).Name()
   140  		dbName = strings.TrimPrefix(dbName, "github.com/cockroachdb/cockroach/pkg/bench.benchmark")
   141  		b.Run(dbName, func(b *testing.B) {
   142  			dbFn(b, fn)
   143  		})
   144  	}
   145  }