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

     1  // Copyright 2021 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 kvbench
    16  
    17  import (
    18  	"fmt"
    19  	"math/rand"
    20  	"os"
    21  	"path"
    22  	"runtime"
    23  	"runtime/pprof"
    24  	"testing"
    25  	"time"
    26  
    27  	"github.com/boltdb/bolt"
    28  
    29  	"github.com/stretchr/testify/require"
    30  )
    31  
    32  const (
    33  	makeProfile = false
    34  )
    35  
    36  // usage: `go test -bench BenchmarkMemoryStore`
    37  func BenchmarkMemoryStore(b *testing.B) {
    38  	benchmarkKVStore(b, newMemStore())
    39  }
    40  
    41  // usage: `go test -bench BenchmarkMemProllyStore`
    42  func BenchmarkMemProllyStore(b *testing.B) {
    43  	benchmarkKVStore(b, newMemoryProllyStore())
    44  }
    45  
    46  // usage: `go test -bench BenchmarkNBSProllyStore`
    47  func BenchmarkNBSProllyStore(b *testing.B) {
    48  	benchmarkKVStore(b, newNBSProllyStore(os.TempDir()))
    49  }
    50  
    51  // usage: `go test -bench BenchmarkBoltStore`
    52  func BenchmarkBoltStore(b *testing.B) {
    53  	benchmarkKVStore(b, newBoltStore(os.TempDir()))
    54  }
    55  
    56  func benchmarkKVStore(b *testing.B, store keyValStore) {
    57  	keys := loadStore(b, store)
    58  
    59  	if makeProfile {
    60  		f := makePprofFile(b)
    61  		err := pprof.StartCPUProfile(f)
    62  		if err != nil {
    63  			b.Fatal(err)
    64  		}
    65  		defer func() {
    66  			pprof.StopCPUProfile()
    67  			if err = f.Close(); err != nil {
    68  				b.Fatal(err)
    69  			}
    70  			fmt.Printf("\twriting CPU profile for %s: %s\n", b.Name(), f.Name())
    71  		}()
    72  	}
    73  
    74  	b.Run("point reads", func(b *testing.B) {
    75  		runBenchmark(b, store, keys)
    76  	})
    77  }
    78  
    79  func loadStore(b *testing.B, store keyValStore) (keys [][]byte) {
    80  	return loadStoreWithParams(b, store, loadParams{
    81  		cardinality: 100_000,
    82  		keySize:     16,
    83  		valSize:     128,
    84  	})
    85  }
    86  
    87  type loadParams struct {
    88  	cardinality uint32
    89  	keySize     uint32
    90  	valSize     uint32
    91  }
    92  
    93  func loadStoreWithParams(b *testing.B, store keyValStore, p loadParams) (keys [][]byte) {
    94  	keys = make([][]byte, 0, p.cardinality)
    95  
    96  	// generate 10K rows at a time
    97  	const batchSize = uint32(10_000)
    98  	numBatches := p.cardinality / batchSize
    99  
   100  	pairSize := p.keySize + p.valSize
   101  	bufSize := pairSize * batchSize
   102  	buf := make([]byte, bufSize)
   103  
   104  	for i := uint32(0); i < numBatches; i++ {
   105  		_, err := rand.Read(buf)
   106  		require.NoError(b, err)
   107  
   108  		kk := make([][]byte, batchSize)
   109  		vv := make([][]byte, batchSize)
   110  
   111  		for j := uint32(0); j < batchSize; j++ {
   112  			offset := j * pairSize
   113  			kk[j] = buf[offset : offset+p.keySize]
   114  			vv[j] = buf[offset+p.keySize : offset+pairSize]
   115  		}
   116  		store.putMany(kk, vv)
   117  		keys = append(keys, kk...)
   118  	}
   119  
   120  	return
   121  }
   122  
   123  func runBenchmark(b *testing.B, store keyValStore, keys [][]byte) {
   124  	runBenchmarkWithParams(b, store, keys, benchParams{})
   125  }
   126  
   127  type benchParams struct{}
   128  
   129  func runBenchmarkWithParams(b *testing.B, store keyValStore, keys [][]byte, p benchParams) {
   130  	if bs, ok := store.(boltStore); ok {
   131  		err := bs.DB.View(func(tx *bolt.Tx) (err error) {
   132  			bk := tx.Bucket([]byte(bucketName))
   133  			err = bk.ForEach(func(k, v []byte) error {
   134  				return nil
   135  			})
   136  			require.NoError(b, err)
   137  			return nil
   138  		})
   139  		require.NoError(b, err)
   140  	}
   141  
   142  	for _, k := range keys {
   143  		_, ok := store.get(k)
   144  		require.True(b, ok)
   145  	}
   146  }
   147  
   148  func makePprofFile(b *testing.B) *os.File {
   149  	_, testFile, _, _ := runtime.Caller(0)
   150  
   151  	name := fmt.Sprintf("%s_%d.pprof", b.Name(), time.Now().Unix())
   152  	f, err := os.Create(path.Join(path.Dir(testFile), name))
   153  	if err != nil {
   154  		b.Fatal(err)
   155  	}
   156  	return f
   157  }