github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/tests/bank_test.go (about)

     1  // Copyright 2015 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 tests_test
    12  
    13  import (
    14  	"bytes"
    15  	"fmt"
    16  	"math/rand"
    17  	"testing"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/bench"
    20  	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
    21  )
    22  
    23  // maxTransfer is the maximum amount to transfer in one transaction.
    24  const maxTransfer = 999
    25  
    26  // runBenchmarkBank mirrors the SQL performed by examples/sql_bank, but
    27  // structured as a benchmark for easier usage of the Go performance analysis
    28  // tools like pprof, memprof and trace.
    29  func runBenchmarkBank(b *testing.B, db *sqlutils.SQLRunner, numAccounts int) {
    30  	{
    31  		// Initialize the "bank" table.
    32  		schema := `
    33  CREATE TABLE IF NOT EXISTS bench.bank (
    34    id INT PRIMARY KEY,
    35    balance INT NOT NULL
    36  )`
    37  		db.Exec(b, schema)
    38  		db.Exec(b, "TRUNCATE TABLE bench.bank")
    39  
    40  		var placeholders bytes.Buffer
    41  		var values []interface{}
    42  		for i := 0; i < numAccounts; i++ {
    43  			if i > 0 {
    44  				placeholders.WriteString(", ")
    45  			}
    46  			fmt.Fprintf(&placeholders, "($%d, 10000)", i+1)
    47  			values = append(values, i)
    48  		}
    49  		stmt := `INSERT INTO bench.bank (id, balance) VALUES ` + placeholders.String()
    50  		db.Exec(b, stmt, values...)
    51  	}
    52  
    53  	b.ResetTimer()
    54  	b.RunParallel(func(pb *testing.PB) {
    55  		for pb.Next() {
    56  			from := rand.Intn(numAccounts)
    57  			to := rand.Intn(numAccounts - 1)
    58  			for from == to {
    59  				to = numAccounts - 1
    60  			}
    61  
    62  			amount := rand.Intn(maxTransfer)
    63  
    64  			const update = `
    65  UPDATE bench.bank
    66    SET balance = CASE id WHEN $1 THEN balance-$3 WHEN $2 THEN balance+$3 END
    67    WHERE id IN ($1, $2) AND (SELECT balance >= $3 FROM bench.bank WHERE id = $1)`
    68  			db.Exec(b, update, from, to, amount)
    69  		}
    70  	})
    71  	b.StopTimer()
    72  }
    73  
    74  func BenchmarkBank(b *testing.B) {
    75  	bench.ForEachDB(b, func(b *testing.B, db *sqlutils.SQLRunner) {
    76  		for _, numAccounts := range []int{2, 4, 8, 32, 64} {
    77  			b.Run(fmt.Sprintf("numAccounts=%d", numAccounts), func(b *testing.B) {
    78  				runBenchmarkBank(b, db, numAccounts)
    79  			})
    80  		}
    81  	})
    82  }