github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/performance/microsysbench/sysbench_test.go (about)

     1  // Copyright 2022 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package microsysbench
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"io"
    21  	"math/rand"
    22  	"os"
    23  	"strconv"
    24  	"strings"
    25  	"testing"
    26  
    27  	"github.com/dolthub/go-mysql-server/sql"
    28  	"github.com/stretchr/testify/require"
    29  
    30  	"github.com/dolthub/dolt/go/cmd/dolt/commands"
    31  	"github.com/dolthub/dolt/go/cmd/dolt/commands/engine"
    32  	"github.com/dolthub/dolt/go/libraries/doltcore/dtestutils"
    33  	"github.com/dolthub/dolt/go/libraries/doltcore/env"
    34  )
    35  
    36  const (
    37  	tableSize   = 10_000
    38  	dataFile    = "testdata.sql"
    39  	createTable = "CREATE TABLE `sbtest1` (" +
    40  		" `id` int NOT NULL AUTO_INCREMENT," +
    41  		" `k` int NOT NULL DEFAULT '0'," +
    42  		" `c` char(120) NOT NULL DEFAULT ''," +
    43  		" `pad` char(60) NOT NULL DEFAULT ''," +
    44  		" PRIMARY KEY (`id`)," +
    45  		" KEY `k_1` (`k`)" +
    46  		");"
    47  )
    48  
    49  var dEnv *env.DoltEnv
    50  
    51  func init() {
    52  	dEnv = dtestutils.CreateTestEnv()
    53  	populateRepo(dEnv, readTestData(dataFile))
    54  }
    55  
    56  func BenchmarkOltpPointSelect(b *testing.B) {
    57  	benchmarkSysbenchQuery(b, func(int) string {
    58  		q := "SELECT c FROM sbtest1 WHERE id=%d"
    59  		return fmt.Sprintf(q, rand.Intn(tableSize))
    60  	})
    61  }
    62  
    63  func BenchmarkOltpJoinScan(b *testing.B) {
    64  	benchmarkSysbenchQuery(b, func(int) string {
    65  		return `select a.id, a.k 
    66  				from sbtest1 a, sbtest1 b 
    67  				where a.id = b.id limit 500`
    68  	})
    69  }
    70  
    71  func BenchmarkProjectionAggregation(b *testing.B) {
    72  	benchmarkSysbenchQuery(b, func(int) string {
    73  		q := "SELECT c, count(id) FROM sbtest1 WHERE k > %d GROUP BY c ORDER BY c"
    74  		return fmt.Sprintf(q, rand.Intn(tableSize))
    75  	})
    76  }
    77  
    78  func BenchmarkSelectRandomPoints(b *testing.B) {
    79  	benchmarkSysbenchQuery(b, func(int) string {
    80  		var sb strings.Builder
    81  		sb.Grow(120)
    82  		sb.WriteString("SELECT id, k, c, pad FROM sbtest1 WHERE k IN (")
    83  		sb.WriteString(strconv.Itoa(rand.Intn(tableSize)))
    84  		for i := 1; i < 10; i++ {
    85  			sb.WriteString(", ")
    86  			sb.WriteString(strconv.Itoa(rand.Intn(tableSize)))
    87  		}
    88  		sb.WriteString(");")
    89  		return sb.String()
    90  	})
    91  }
    92  
    93  func benchmarkSysbenchQuery(b *testing.B, getQuery func(int) string) {
    94  	ctx, eng := setupBenchmark(b, dEnv)
    95  	for i := 0; i < b.N; i++ {
    96  		_, iter, err := eng.Query(ctx, getQuery(i))
    97  		require.NoError(b, err)
    98  		for {
    99  			if _, err = iter.Next(ctx); err != nil {
   100  				break
   101  			}
   102  		}
   103  		require.Error(b, io.EOF)
   104  		err = iter.Close(ctx)
   105  		require.NoError(b, err)
   106  	}
   107  	_ = eng.Close()
   108  	b.ReportAllocs()
   109  }
   110  
   111  func setupBenchmark(t *testing.B, dEnv *env.DoltEnv) (*sql.Context, *engine.SqlEngine) {
   112  	ctx := context.Background()
   113  	config := &engine.SqlEngineConfig{
   114  		ServerUser: "root",
   115  		Autocommit: true,
   116  	}
   117  
   118  	mrEnv, err := env.MultiEnvForDirectory(ctx, dEnv.Config.WriteableConfig(), dEnv.FS, dEnv.Version, dEnv)
   119  	require.NoError(t, err)
   120  
   121  	eng, err := engine.NewSqlEngine(ctx, mrEnv, config)
   122  	require.NoError(t, err)
   123  
   124  	sqlCtx, err := eng.NewLocalContext(ctx)
   125  	require.NoError(t, err)
   126  
   127  	sqlCtx.SetCurrentDatabase("dolt")
   128  	return sqlCtx, eng
   129  }
   130  
   131  func readTestData(file string) string {
   132  	data, err := os.ReadFile(file)
   133  	if err != nil {
   134  		panic(err)
   135  	}
   136  	return string(data)
   137  }
   138  
   139  func populateRepo(dEnv *env.DoltEnv, insertData string) {
   140  	execSql := func(dEnv *env.DoltEnv, q string) int {
   141  		ctx := context.Background()
   142  		args := []string{"-r", "null", "-q", q}
   143  		cliCtx, err := commands.NewArgFreeCliContext(ctx, dEnv)
   144  		if err != nil {
   145  			panic(err)
   146  		}
   147  
   148  		return commands.SqlCmd{}.Exec(ctx, "sql", args, dEnv, cliCtx)
   149  	}
   150  	execSql(dEnv, createTable)
   151  	execSql(dEnv, insertData)
   152  }