func BenchmarkSstRekey(b *testing.B) {
// TODO(dan): DRY this with BenchmarkRocksDBSstFileReader.
dir, cleanupFn := testutils.TempDir(b, 1)
defer cleanupFn()
sstPath := filepath.Join(dir, "sst")
{
const maxEntries = 100000
const keyLen = 10
const valLen = 100
b.SetBytes(keyLen + valLen)
ts := hlc.Timestamp{WallTime: timeutil.Now().UnixNano()}
kv := engine.MVCCKeyValue{
Key: engine.MVCCKey{Key: roachpb.Key(make([]byte, keyLen)), Timestamp: ts},
Value: make([]byte, valLen),
}
sst := engine.MakeRocksDBSstFileWriter()
if err := sst.Open(sstPath); err != nil {
b.Fatal(sst)
}
var entries = b.N
if entries > maxEntries {
entries = maxEntries
}
for i := 0; i < entries; i++ {
payload := []byte(fmt.Sprintf("%09d", i))
kv.Key.Key = kv.Key.Key[:0]
kv.Key.Key = encoding.EncodeUvarintAscending(kv.Key.Key, uint64(i)) // tableID
kv.Key.Key = encoding.EncodeUvarintAscending(kv.Key.Key, 0) // indexID
kv.Key.Key = encoding.EncodeBytesAscending(kv.Key.Key, payload)
kv.Key.Key = keys.MakeRowSentinelKey(kv.Key.Key)
copy(kv.Value, payload)
if err := sst.Add(kv); err != nil {
b.Fatal(err)
}
}
if err := sst.Close(); err != nil {
b.Fatal(err)
}
}
const newTableID = 100
b.ResetTimer()
sst, err := engine.MakeRocksDBSstFileReader()
if err != nil {
b.Fatal(err)
}
if err := sst.AddFile(sstPath); err != nil {
b.Fatal(err)
}
defer sst.Close()
count := 0
iterateFn := sql.MakeRekeyMVCCKeyValFunc(newTableID, func(kv engine.MVCCKeyValue) (bool, error) {
count++
if count >= b.N {
return true, nil
}
return false, nil
})
for {
if err := sst.Iterate(engine.MVCCKey{Key: keys.MinKey}, engine.MVCCKey{Key: keys.MaxKey}, iterateFn); err != nil {
b.Fatal(err)
}
if count >= b.N {
break
}
}
}
// Ingest loads some data in an sstable into an empty range. Only the keys
// between startKey and endKey are loaded. If newTableID is non-zero, every
// row's key is rewritten to be for that table.
func Ingest(
ctx context.Context,
txn *client.Txn,
path string,
checksum uint32,
startKey, endKey roachpb.Key,
newTableID sqlbase.ID,
) error {
// TODO(mjibson): An appropriate value for this should be determined. The
// current value was guessed at but appears to work well.
const batchSize = 10000
// TODO(dan): Check if the range being ingested into is empty. If newTableID
// is non-zero, it'll have to be derived from startKey and endKey.
f, err := os.Open(path)
if err != nil {
return err
}
defer f.Close()
crc := crc32.New(crc32.MakeTable(crc32.Castagnoli))
if _, err := io.Copy(crc, f); err != nil {
return nil
}
if c := crc.Sum32(); c != checksum {
return errors.Errorf("%s: checksum mismatch got %d expected %d", path, c, checksum)
}
sst, err := engine.MakeRocksDBSstFileReader()
if err != nil {
return err
}
defer sst.Close()
if err := sst.AddFile(path); err != nil {
return err
}
b := txn.NewBatch()
var v roachpb.Value
count := 0
ingestFunc := func(kv engine.MVCCKeyValue) (bool, error) {
v = roachpb.Value{RawBytes: kv.Value}
v.ClearChecksum()
if log.V(3) {
log.Infof(ctx, "Put %s %s\n", kv.Key.Key, v.PrettyPrint())
}
b.Put(kv.Key.Key, &v)
count++
if count > batchSize {
if err := txn.Run(b); err != nil {
return true, err
}
b = txn.NewBatch()
count = 0
}
return false, nil
}
if newTableID != 0 {
// MakeRekeyMVCCKeyValFunc modifies the keys, but this is safe because
// the one we get back from rocksDBIterator.Key is a copy (not a
// reference to the mmaped file.)
ingestFunc = MakeRekeyMVCCKeyValFunc(newTableID, ingestFunc)
}
startKeyMVCC, endKeyMVCC := engine.MVCCKey{Key: startKey}, engine.MVCCKey{Key: endKey}
if err := sst.Iterate(startKeyMVCC, endKeyMVCC, ingestFunc); err != nil {
return err
}
return txn.Run(b)
}