github.com/rosedblabs/rosedb/v2@v2.3.7-0.20240423093736-a89ea823e5b9/benchmark/bench_test.go (about)

     1  package benchmark
     2  
     3  import (
     4  	"errors"
     5  	"math/rand"
     6  	"os"
     7  	"testing"
     8  
     9  	"github.com/rosedblabs/rosedb/v2"
    10  	"github.com/rosedblabs/rosedb/v2/utils"
    11  	"github.com/stretchr/testify/assert"
    12  )
    13  
    14  var db *rosedb.DB
    15  
    16  func openDB() func() {
    17  	options := rosedb.DefaultOptions
    18  	options.DirPath = "/tmp/rosedbv2"
    19  
    20  	var err error
    21  	db, err = rosedb.Open(options)
    22  	if err != nil {
    23  		panic(err)
    24  	}
    25  
    26  	return func() {
    27  		_ = db.Close()
    28  		_ = os.RemoveAll(options.DirPath)
    29  	}
    30  }
    31  
    32  func BenchmarkPutGet(b *testing.B) {
    33  	closer := openDB()
    34  	defer closer()
    35  
    36  	b.Run("put", benchmarkPut)
    37  	b.Run("get", bencharkGet)
    38  }
    39  
    40  func BenchmarkBatchPutGet(b *testing.B) {
    41  	closer := openDB()
    42  	defer closer()
    43  
    44  	b.Run("batchPut", benchmarkBatchPut)
    45  	b.Run("batchGet", benchmarkBatchGet)
    46  }
    47  
    48  func benchmarkPut(b *testing.B) {
    49  	b.ResetTimer()
    50  	b.ReportAllocs()
    51  
    52  	for i := 0; i < b.N; i++ {
    53  		err := db.Put(utils.GetTestKey(i), utils.RandomValue(1024))
    54  		assert.Nil(b, err)
    55  	}
    56  }
    57  
    58  func benchmarkBatchPut(b *testing.B) {
    59  	b.ResetTimer()
    60  	b.ReportAllocs()
    61  
    62  	batch := db.NewBatch(rosedb.DefaultBatchOptions)
    63  	defer batch.Commit()
    64  	for i := 0; i < b.N; i++ {
    65  		err := batch.Put(utils.GetTestKey(i), utils.RandomValue(1024))
    66  		assert.Nil(b, err)
    67  	}
    68  }
    69  
    70  func benchmarkBatchGet(b *testing.B) {
    71  	for i := 0; i < 10000; i++ {
    72  		err := db.Put(utils.GetTestKey(i), utils.RandomValue(1024))
    73  		assert.Nil(b, err)
    74  	}
    75  
    76  	b.ResetTimer()
    77  	b.ReportAllocs()
    78  	batch := db.NewBatch(rosedb.DefaultBatchOptions)
    79  	defer batch.Commit()
    80  	for i := 0; i < b.N; i++ {
    81  		_, err := batch.Get(utils.GetTestKey(rand.Int()))
    82  		if err != nil && !errors.Is(err, rosedb.ErrKeyNotFound) {
    83  			b.Fatal(err)
    84  		}
    85  	}
    86  }
    87  
    88  func bencharkGet(b *testing.B) {
    89  	for i := 0; i < 10000; i++ {
    90  		err := db.Put(utils.GetTestKey(i), utils.RandomValue(1024))
    91  		assert.Nil(b, err)
    92  	}
    93  
    94  	b.ResetTimer()
    95  	b.ReportAllocs()
    96  
    97  	for i := 0; i < b.N; i++ {
    98  		_, err := db.Get(utils.GetTestKey(rand.Int()))
    99  		if err != nil && !errors.Is(err, rosedb.ErrKeyNotFound) {
   100  			b.Fatal(err)
   101  		}
   102  	}
   103  }