kythe.io@v0.0.68-0.20240422202219-7225dbc01741/kythe/go/test/storage/keyvalue/keyvalue.go (about)

     1  /*
     2   * Copyright 2015 The Kythe Authors. All rights reserved.
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *   http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  // Package keyvalue contains utilities to test keyvalue DB implementations.
    18  package keyvalue // import "kythe.io/kythe/go/test/storage/keyvalue"
    19  
    20  import (
    21  	"context"
    22  	"testing"
    23  
    24  	"kythe.io/kythe/go/storage/keyvalue"
    25  	"kythe.io/kythe/go/test/testutil"
    26  )
    27  
    28  // DB re-exports keyvalue.DB for tests
    29  type DB keyvalue.DB
    30  
    31  // NewGraphStore re-exports keyvalue.NewGraphStore for tests
    32  var NewGraphStore = keyvalue.NewGraphStore
    33  
    34  // CreateFunc creates a temporary keyvalue.DB with a corresponding function to
    35  // destroy it completely.
    36  type CreateFunc func() (DB, DestroyFunc, error)
    37  
    38  // DestroyFunc destroys its corresponding keyvalue.DB returned from a
    39  // CreateFunc.
    40  type DestroyFunc func() error
    41  
    42  // NullDestroy does nothing.
    43  func NullDestroy() error { return nil }
    44  
    45  const (
    46  	keySize = 16 // bytes
    47  	valSize = 32 // bytes
    48  )
    49  
    50  var ctx = context.Background()
    51  
    52  // BatchWriteBenchmark benchmarks the Write method of the given keyvalue.DB.
    53  // The number of updates per write is configured with batchSize.
    54  func BatchWriteBenchmark(b *testing.B, create CreateFunc, batchSize int) {
    55  	db, destroy, err := create()
    56  	testutil.Fatalf(b, "CreateFunc error: %v", err)
    57  	defer func() {
    58  		testutil.Fatalf(b, "db close error: %v", db.Close(ctx))
    59  		testutil.Fatalf(b, "DestroyFunc error: %v", destroy())
    60  	}()
    61  
    62  	keyBuf := make([]byte, keySize)
    63  	valBuf := make([]byte, valSize)
    64  	for i := 0; i < b.N; i++ {
    65  		wr, err := db.Writer(ctx)
    66  		testutil.Fatalf(b, "writer error: %v", err)
    67  		for j := 0; j < batchSize; j++ {
    68  			testutil.RandBytes(keyBuf)
    69  			testutil.RandBytes(valBuf)
    70  			testutil.Fatalf(b, "write error: %v", wr.Write(keyBuf, valBuf))
    71  		}
    72  		testutil.Fatalf(b, "writer close error: %v", wr.Close())
    73  	}
    74  }
    75  
    76  // BatchWriteParallelBenchmark benchmarks the Write method of the given
    77  // keyvalue.DB in parallel.  The number of updates per write is configured with
    78  // batchSize.
    79  func BatchWriteParallelBenchmark(b *testing.B, create CreateFunc, batchSize int) {
    80  	db, destroy, err := create()
    81  	testutil.Fatalf(b, "CreateFunc error: %v", err)
    82  	defer func() {
    83  		testutil.Fatalf(b, "db close error: %v", db.Close(ctx))
    84  		testutil.Fatalf(b, "DestroyFunc error: %v", destroy())
    85  	}()
    86  
    87  	b.RunParallel(func(pb *testing.PB) {
    88  		keyBuf := make([]byte, keySize)
    89  		valBuf := make([]byte, valSize)
    90  		for pb.Next() {
    91  			wr, err := db.Writer(ctx)
    92  			testutil.Fatalf(b, "writer error: %v", err)
    93  			for j := 0; j < batchSize; j++ {
    94  				testutil.RandBytes(keyBuf)
    95  				testutil.RandBytes(valBuf)
    96  				testutil.Fatalf(b, "write error: %v", wr.Write(keyBuf, valBuf))
    97  			}
    98  			testutil.Fatalf(b, "writer close error: %v", wr.Close())
    99  		}
   100  	})
   101  }