func maybeUnmarshalInline(v []byte, dest proto.Message) error {
var meta enginepb.MVCCMetadata
if err := meta.Unmarshal(v); err != nil {
return err
}
value := roachpb.Value{
RawBytes: meta.RawBytes,
}
return value.GetProto(dest)
}
func copySeqCache(
e engine.ReadWriter,
ms *enginepb.MVCCStats,
srcID, dstID roachpb.RangeID,
keyMin, keyMax engine.MVCCKey,
) (int, error) {
var scratch [64]byte
var count int
var meta enginepb.MVCCMetadata
// TODO(spencer): look into making this an MVCCIteration and writing
// the values using MVCC so we can avoid the ugliness of updating
// the MVCCStats by hand below.
err := e.Iterate(keyMin, keyMax,
func(kv engine.MVCCKeyValue) (bool, error) {
// Decode the key, skipping on error. Otherwise, write it to the
// corresponding key in the new cache.
txnID, err := decodeAbortCacheMVCCKey(kv.Key, scratch[:0])
if err != nil {
return false, errors.Errorf("could not decode an abort cache key %s: %s", kv.Key, err)
}
key := keys.AbortCacheKey(dstID, txnID)
encKey := engine.MakeMVCCMetadataKey(key)
// Decode the MVCCMetadata value.
if err := proto.Unmarshal(kv.Value, &meta); err != nil {
return false, errors.Errorf("could not decode mvcc metadata %s [% x]: %s", kv.Key, kv.Value, err)
}
value := engine.MakeValue(meta)
value.ClearChecksum()
value.InitChecksum(key)
meta.RawBytes = value.RawBytes
keyBytes, valBytes, err := engine.PutProto(e, encKey, &meta)
if err != nil {
return false, err
}
count++
if ms != nil {
ms.SysBytes += keyBytes + valBytes
ms.SysCount++
}
return false, nil
})
return count, err
}
func tryRangeIDKey(kv engine.MVCCKeyValue) (string, error) {
if kv.Key.Timestamp != hlc.ZeroTimestamp {
return "", fmt.Errorf("range ID keys shouldn't have timestamps: %s", kv.Key)
}
_, _, suffix, _, err := keys.DecodeRangeIDKey(kv.Key.Key)
if err != nil {
return "", err
}
// All range ID keys are stored inline on the metadata.
var meta enginepb.MVCCMetadata
if err := meta.Unmarshal(kv.Value); err != nil {
return "", err
}
value := roachpb.Value{RawBytes: meta.RawBytes}
// Values encoded as protobufs set msg and continue outside the
// switch. Other types are handled inside the switch and return.
var msg proto.Message
switch {
case bytes.Equal(suffix, keys.LocalLeaseAppliedIndexSuffix):
fallthrough
case bytes.Equal(suffix, keys.LocalRaftAppliedIndexSuffix):
i, err := value.GetInt()
if err != nil {
return "", err
}
return strconv.FormatInt(i, 10), nil
case bytes.Equal(suffix, keys.LocalRangeFrozenStatusSuffix):
b, err := value.GetBool()
if err != nil {
return "", err
}
return strconv.FormatBool(b), nil
case bytes.Equal(suffix, keys.LocalAbortCacheSuffix):
msg = &roachpb.AbortCacheEntry{}
case bytes.Equal(suffix, keys.LocalRangeLastGCSuffix):
msg = &hlc.Timestamp{}
case bytes.Equal(suffix, keys.LocalRaftTombstoneSuffix):
msg = &roachpb.RaftTombstone{}
case bytes.Equal(suffix, keys.LocalRaftTruncatedStateSuffix):
msg = &roachpb.RaftTruncatedState{}
case bytes.Equal(suffix, keys.LocalRangeLeaseSuffix):
msg = &roachpb.Lease{}
case bytes.Equal(suffix, keys.LocalRangeStatsSuffix):
msg = &enginepb.MVCCStats{}
case bytes.Equal(suffix, keys.LocalRaftHardStateSuffix):
msg = &raftpb.HardState{}
case bytes.Equal(suffix, keys.LocalRaftLastIndexSuffix):
i, err := value.GetInt()
if err != nil {
return "", err
}
return strconv.FormatInt(i, 10), nil
case bytes.Equal(suffix, keys.LocalRangeLastVerificationTimestampSuffix):
msg = &hlc.Timestamp{}
case bytes.Equal(suffix, keys.LocalRangeLastReplicaGCTimestampSuffix):
msg = &hlc.Timestamp{}
default:
return "", fmt.Errorf("unknown raft id key %s", suffix)
}
if err := value.GetProto(msg); err != nil {
return "", err
}
return msg.String(), nil
}
func TestEngineBatch(t *testing.T) {
defer leaktest.AfterTest(t)()
runWithAllEngines(func(engine Engine, t *testing.T) {
numShuffles := 100
key := mvccKey("a")
// Those are randomized below.
type data struct {
key MVCCKey
value []byte
merge bool
}
batch := []data{
{key, appender("~ockroachDB"), false},
{key, appender("C~ckroachDB"), false},
{key, appender("Co~kroachDB"), false},
{key, appender("Coc~roachDB"), false},
{key, appender("C**k~oachDB"), false},
{key, appender("Cockr~achDB"), false},
{key, appender("Cockro~chDB"), false},
{key, appender("Cockroa~hDB"), false},
{key, appender("Cockroac~DB"), false},
{key, appender("Cockroach~B"), false},
{key, appender("CockroachD~"), false},
{key, nil, false},
{key, appender("C"), true},
{key, appender(" o"), true},
{key, appender(" c"), true},
{key, appender(" k"), true},
{key, appender("r"), true},
{key, appender(" o"), true},
{key, appender(" a"), true},
{key, appender(" c"), true},
{key, appender("h"), true},
{key, appender(" D"), true},
{key, appender(" B"), true},
}
apply := func(eng ReadWriter, d data) error {
if d.value == nil {
return eng.Clear(d.key)
} else if d.merge {
return eng.Merge(d.key, d.value)
}
return eng.Put(d.key, d.value)
}
get := func(eng ReadWriter, key MVCCKey) []byte {
b, err := eng.Get(key)
if err != nil {
t.Fatal(err)
}
var m enginepb.MVCCMetadata
if err := proto.Unmarshal(b, &m); err != nil {
t.Fatal(err)
}
if !m.IsInline() {
return nil
}
valueBytes, err := MakeValue(m).GetBytes()
if err != nil {
t.Fatal(err)
}
return valueBytes
}
for i := 0; i < numShuffles; i++ {
// In each run, create an array of shuffled operations.
shuffledIndices := rand.Perm(len(batch))
currentBatch := make([]data, len(batch))
for k := range currentBatch {
currentBatch[k] = batch[shuffledIndices[k]]
}
// Reset the key
if err := engine.Clear(key); err != nil {
t.Fatal(err)
}
// Run it once with individual operations and remember the result.
for i, op := range currentBatch {
if err := apply(engine, op); err != nil {
t.Errorf("%d: op %v: %v", i, op, err)
continue
}
}
expectedValue := get(engine, key)
// Run the whole thing as a batch and compare.
b := engine.NewBatch()
if err := b.Clear(key); err != nil {
t.Fatal(err)
}
for _, op := range currentBatch {
if err := apply(b, op); err != nil {
t.Fatal(err)
}
}
// Try getting the value from the batch.
actualValue := get(b, key)
if !bytes.Equal(actualValue, expectedValue) {
t.Errorf("%d: expected %s, but got %s", i, expectedValue, actualValue)
}
// Try using an iterator to get the value from the batch.
iter := b.NewIterator(false)
iter.Seek(key)
if !iter.Valid() {
if currentBatch[len(currentBatch)-1].value != nil {
t.Errorf("%d: batch seek invalid", i)
}
} else if !iter.Key().Equal(key) {
t.Errorf("%d: batch seek expected key %s, but got %s", i, key, iter.Key())
} else {
var m enginepb.MVCCMetadata
if err := iter.ValueProto(&m); err != nil {
t.Fatal(err)
}
valueBytes, err := MakeValue(m).GetBytes()
if err != nil {
t.Fatal(err)
}
if !bytes.Equal(valueBytes, expectedValue) {
t.Errorf("%d: expected %s, but got %s", i, expectedValue, valueBytes)
}
}
iter.Close()
// Commit the batch and try getting the value from the engine.
if err := b.Commit(); err != nil {
t.Errorf("%d: %v", i, err)
continue
}
actualValue = get(engine, key)
if !bytes.Equal(actualValue, expectedValue) {
t.Errorf("%d: expected %s, but got %s", i, expectedValue, actualValue)
}
}
}, t)
}