github.com/ava-labs/avalanchego@v1.11.11/database/dbtest/benchmark.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package dbtest
     5  
     6  import (
     7  	"math/rand"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/require"
    11  
    12  	"github.com/ava-labs/avalanchego/database"
    13  	"github.com/ava-labs/avalanchego/utils/units"
    14  )
    15  
    16  var (
    17  	// Benchmarks is a list of all database benchmarks
    18  	Benchmarks = map[string]func(b *testing.B, db database.Database, keys, values [][]byte){
    19  		"Get":            BenchmarkGet,
    20  		"Put":            BenchmarkPut,
    21  		"Delete":         BenchmarkDelete,
    22  		"BatchPut":       BenchmarkBatchPut,
    23  		"BatchDelete":    BenchmarkBatchDelete,
    24  		"BatchWrite":     BenchmarkBatchWrite,
    25  		"ParallelGet":    BenchmarkParallelGet,
    26  		"ParallelPut":    BenchmarkParallelPut,
    27  		"ParallelDelete": BenchmarkParallelDelete,
    28  	}
    29  	// BenchmarkSizes to use with each benchmark
    30  	BenchmarkSizes = [][]int{
    31  		// count, keySize, valueSize
    32  		{1024, 32, 32},
    33  		{1024, 256, 256},
    34  		{1024, 2 * units.KiB, 2 * units.KiB},
    35  	}
    36  )
    37  
    38  // Writes size data into the db in order to setup reads in subsequent tests.
    39  func SetupBenchmark(b *testing.B, count int, keySize, valueSize int) ([][]byte, [][]byte) {
    40  	require := require.New(b)
    41  
    42  	b.Helper()
    43  
    44  	keys := make([][]byte, count)
    45  	values := make([][]byte, count)
    46  	for i := 0; i < count; i++ {
    47  		keyBytes := make([]byte, keySize)
    48  		valueBytes := make([]byte, valueSize)
    49  		_, err := rand.Read(keyBytes) // #nosec G404
    50  		require.NoError(err)
    51  		_, err = rand.Read(valueBytes) // #nosec G404
    52  		require.NoError(err)
    53  		keys[i], values[i] = keyBytes, valueBytes
    54  	}
    55  	return keys, values
    56  }
    57  
    58  // BenchmarkGet measures the time it takes to get an operation from a database.
    59  func BenchmarkGet(b *testing.B, db database.Database, keys, values [][]byte) {
    60  	require.NotEmpty(b, keys)
    61  	count := len(keys)
    62  
    63  	require := require.New(b)
    64  
    65  	for i, key := range keys {
    66  		value := values[i]
    67  		require.NoError(db.Put(key, value))
    68  	}
    69  
    70  	b.ResetTimer()
    71  
    72  	// Reads b.N values from the db
    73  	for i := 0; i < b.N; i++ {
    74  		_, err := db.Get(keys[i%count])
    75  		require.NoError(err)
    76  	}
    77  }
    78  
    79  // BenchmarkPut measures the time it takes to write an operation to a database.
    80  func BenchmarkPut(b *testing.B, db database.Database, keys, values [][]byte) {
    81  	require.NotEmpty(b, keys)
    82  	count := len(keys)
    83  
    84  	// Writes b.N values to the db
    85  	for i := 0; i < b.N; i++ {
    86  		require.NoError(b, db.Put(keys[i%count], values[i%count]))
    87  	}
    88  }
    89  
    90  // BenchmarkDelete measures the time it takes to delete a (k, v) from a database.
    91  func BenchmarkDelete(b *testing.B, db database.Database, keys, values [][]byte) {
    92  	require.NotEmpty(b, keys)
    93  	count := len(keys)
    94  
    95  	require := require.New(b)
    96  
    97  	// Writes random values of size _size_ to the database
    98  	for i, key := range keys {
    99  		value := values[i]
   100  		require.NoError(db.Put(key, value))
   101  	}
   102  
   103  	b.ResetTimer()
   104  
   105  	// Deletes b.N values from the db
   106  	for i := 0; i < b.N; i++ {
   107  		require.NoError(db.Delete(keys[i%count]))
   108  	}
   109  }
   110  
   111  // BenchmarkBatchPut measures the time it takes to batch put.
   112  func BenchmarkBatchPut(b *testing.B, db database.Database, keys, values [][]byte) {
   113  	require.NotEmpty(b, keys)
   114  	count := len(keys)
   115  
   116  	batch := db.NewBatch()
   117  	for i := 0; i < b.N; i++ {
   118  		require.NoError(b, batch.Put(keys[i%count], values[i%count]))
   119  	}
   120  }
   121  
   122  // BenchmarkBatchDelete measures the time it takes to batch delete.
   123  func BenchmarkBatchDelete(b *testing.B, db database.Database, keys, _ [][]byte) {
   124  	require.NotEmpty(b, keys)
   125  	count := len(keys)
   126  
   127  	batch := db.NewBatch()
   128  	for i := 0; i < b.N; i++ {
   129  		require.NoError(b, batch.Delete(keys[i%count]))
   130  	}
   131  }
   132  
   133  // BenchmarkBatchWrite measures the time it takes to batch write.
   134  func BenchmarkBatchWrite(b *testing.B, db database.Database, keys, values [][]byte) {
   135  	require.NotEmpty(b, keys)
   136  
   137  	require := require.New(b)
   138  
   139  	batch := db.NewBatch()
   140  	for i, key := range keys {
   141  		value := values[i]
   142  		require.NoError(batch.Put(key, value))
   143  	}
   144  
   145  	b.ResetTimer()
   146  
   147  	for i := 0; i < b.N; i++ {
   148  		require.NoError(batch.Write())
   149  	}
   150  }
   151  
   152  // BenchmarkParallelGet measures the time it takes to read in parallel.
   153  func BenchmarkParallelGet(b *testing.B, db database.Database, keys, values [][]byte) {
   154  	require.NotEmpty(b, keys)
   155  	count := len(keys)
   156  
   157  	require := require.New(b)
   158  
   159  	for i, key := range keys {
   160  		value := values[i]
   161  		require.NoError(db.Put(key, value))
   162  	}
   163  
   164  	b.ResetTimer()
   165  
   166  	b.RunParallel(func(pb *testing.PB) {
   167  		for i := 0; pb.Next(); i++ {
   168  			_, err := db.Get(keys[i%count])
   169  			require.NoError(err)
   170  		}
   171  	})
   172  }
   173  
   174  // BenchmarkParallelPut measures the time it takes to write to the db in parallel.
   175  func BenchmarkParallelPut(b *testing.B, db database.Database, keys, values [][]byte) {
   176  	require.NotEmpty(b, keys)
   177  	count := len(keys)
   178  
   179  	b.RunParallel(func(pb *testing.PB) {
   180  		// Write N values to the db
   181  		for i := 0; pb.Next(); i++ {
   182  			require.NoError(b, db.Put(keys[i%count], values[i%count]))
   183  		}
   184  	})
   185  }
   186  
   187  // BenchmarkParallelDelete measures the time it takes to delete a (k, v) from the db.
   188  func BenchmarkParallelDelete(b *testing.B, db database.Database, keys, values [][]byte) {
   189  	require.NotEmpty(b, keys)
   190  	count := len(keys)
   191  
   192  	require := require.New(b)
   193  	for i, key := range keys {
   194  		value := values[i]
   195  		require.NoError(db.Put(key, value))
   196  	}
   197  	b.ResetTimer()
   198  
   199  	b.RunParallel(func(pb *testing.PB) {
   200  		// Deletes b.N values from the db
   201  		for i := 0; pb.Next(); i++ {
   202  			require.NoError(db.Delete(keys[i%count]))
   203  		}
   204  	})
   205  }