github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvnemesis/engine.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package kvnemesis
    12  
    13  import (
    14  	"fmt"
    15  	"strings"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    18  	"github.com/cockroachdb/cockroach/pkg/storage"
    19  	"github.com/cockroachdb/cockroach/pkg/util/bufalloc"
    20  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    21  	"github.com/cockroachdb/pebble"
    22  	"github.com/cockroachdb/pebble/vfs"
    23  )
    24  
    25  // Engine is a simplified version of storage.ReadWriter. It is a multi-version
    26  // key-value map, meaning that each read or write has an associated timestamp
    27  // and a read returns the write for the key with the highest timestamp (which is
    28  // not necessarily the most recently ingested write). Engine is not threadsafe.
    29  type Engine struct {
    30  	kvs *pebble.DB
    31  	b   bufalloc.ByteAllocator
    32  }
    33  
    34  // MakeEngine returns a new Engine.
    35  func MakeEngine() (*Engine, error) {
    36  	opts := storage.DefaultPebbleOptions()
    37  	opts.FS = vfs.NewMem()
    38  	kvs, err := pebble.Open(`kvnemesis`, opts)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  	return &Engine{kvs: kvs}, nil
    43  }
    44  
    45  // Close closes the Engine, freeing associated resources.
    46  func (e *Engine) Close() {
    47  	if err := e.kvs.Close(); err != nil {
    48  		panic(err)
    49  	}
    50  }
    51  
    52  // Get returns the value for this key with the highest timestamp <= ts. If no
    53  // such value exists, the returned value's RawBytes is nil.
    54  func (e *Engine) Get(key roachpb.Key, ts hlc.Timestamp) roachpb.Value {
    55  	iter := e.kvs.NewIter(nil)
    56  	defer func() { _ = iter.Close() }()
    57  	iter.SeekGE(storage.EncodeKey(storage.MVCCKey{Key: key, Timestamp: ts}))
    58  	if !iter.Valid() {
    59  		return roachpb.Value{}
    60  	}
    61  	// This use of iter.Key() is safe because it comes entirely before the
    62  	// deferred iter.Close.
    63  	mvccKey, err := storage.DecodeMVCCKey(iter.Key())
    64  	if err != nil {
    65  		panic(err)
    66  	}
    67  	if !mvccKey.Key.Equal(key) {
    68  		return roachpb.Value{}
    69  	}
    70  	var valCopy []byte
    71  	valCopy, e.b = e.b.Copy(iter.Value(), 0 /* extraCap */)
    72  	return roachpb.Value{RawBytes: valCopy, Timestamp: mvccKey.Timestamp}
    73  }
    74  
    75  // Put inserts a key/value/timestamp tuple. If an exact key/timestamp pair is
    76  // Put again, it overwrites the previous value.
    77  func (e *Engine) Put(key storage.MVCCKey, value []byte) {
    78  	if err := e.kvs.Set(storage.EncodeKey(key), value, nil); err != nil {
    79  		panic(err)
    80  	}
    81  }
    82  
    83  // Iterate calls the given closure with every KV in the Engine, in ascending
    84  // order.
    85  func (e *Engine) Iterate(fn func(key storage.MVCCKey, value []byte, err error)) {
    86  	iter := e.kvs.NewIter(nil)
    87  	defer func() { _ = iter.Close() }()
    88  	for iter.First(); iter.Valid(); iter.Next() {
    89  		if err := iter.Error(); err != nil {
    90  			fn(storage.MVCCKey{}, nil, err)
    91  			continue
    92  		}
    93  		var keyCopy, valCopy []byte
    94  		keyCopy, e.b = e.b.Copy(iter.Key(), 0 /* extraCap */)
    95  		valCopy, e.b = e.b.Copy(iter.Value(), 0 /* extraCap */)
    96  		key, err := storage.DecodeMVCCKey(keyCopy)
    97  		if err != nil {
    98  			fn(storage.MVCCKey{}, nil, err)
    99  			continue
   100  		}
   101  		fn(key, valCopy, nil)
   102  	}
   103  }
   104  
   105  // DebugPrint returns the entire contents of this Engine as a string for use in
   106  // debugging.
   107  func (e *Engine) DebugPrint(indent string) string {
   108  	var buf strings.Builder
   109  	e.Iterate(func(key storage.MVCCKey, value []byte, err error) {
   110  		if buf.Len() > 0 {
   111  			buf.WriteString("\n")
   112  		}
   113  		if err != nil {
   114  			fmt.Fprintf(&buf, "(err:%s)", err)
   115  		} else {
   116  			fmt.Fprintf(&buf, "%s%s %s -> %s",
   117  				indent, key.Key, key.Timestamp, roachpb.Value{RawBytes: value}.PrettyPrint())
   118  		}
   119  	})
   120  	return buf.String()
   121  }