func backupRestoreTestSetup(
t testing.TB, numAccounts int,
) (
ctx context.Context,
tempDir string,
tc *testcluster.TestCluster,
kvDB *client.DB,
sqlDB *sqlutils.SQLRunner,
cleanup func(),
) {
ctx = context.Background()
dir, dirCleanupFn := testutils.TempDir(t, 1)
// Use ReplicationManual so we can force full replication, which is needed
// to later move the leases around.
tc = testcluster.StartTestCluster(t, backupRestoreClusterSize, base.TestClusterArgs{
ReplicationMode: base.ReplicationManual,
})
sqlDB = sqlutils.MakeSQLRunner(t, tc.Conns[0])
kvDB = tc.Server(0).KVClient().(*client.DB)
sqlDB.Exec(bankCreateDatabase)
sqlDB.Exec(bankCreateTable)
for _, insert := range bankDataInsertStmts(numAccounts) {
sqlDB.Exec(insert)
}
for _, split := range bankSplitStmts(numAccounts, backupRestoreDefaultRanges) {
sqlDB.Exec(split)
}
targets := make([]testcluster.ReplicationTarget, backupRestoreClusterSize-1)
for i := 1; i < backupRestoreClusterSize; i++ {
targets[i-1] = tc.Target(i)
}
txn := client.NewTxn(ctx, *kvDB)
rangeDescs, err := sql.AllRangeDescriptors(txn)
if err != nil {
t.Fatal(err)
}
for _, r := range rangeDescs {
if _, err := tc.AddReplicas(r.StartKey.AsRawKey(), targets...); err != nil {
t.Fatal(err)
}
}
cleanupFn := func() {
tc.Stopper().Stop()
dirCleanupFn()
}
return ctx, dir, tc, kvDB, sqlDB, cleanupFn
}
func TestRocksDBTimeBound(t *testing.T) {
defer leaktest.AfterTest(t)()
dir, dirCleanup := testutils.TempDir(t, 0)
defer dirCleanup()
rocksdb, err := NewRocksDB(roachpb.Attributes{}, dir, RocksDBCache{}, 0, DefaultMaxOpenFiles)
if err != nil {
t.Fatalf("could not create new rocksdb db instance at %s: %v", dir, err)
}
defer rocksdb.Close()
var minTimestamp = hlc.Timestamp{WallTime: 1, Logical: 0}
var maxTimestamp = hlc.Timestamp{WallTime: 3, Logical: 0}
times := []hlc.Timestamp{
{WallTime: 2, Logical: 0},
minTimestamp,
maxTimestamp,
{WallTime: 2, Logical: 0},
}
for i, time := range times {
s := fmt.Sprintf("%02d", i)
key := MVCCKey{Key: roachpb.Key(s), Timestamp: time}
if err := rocksdb.Put(key, []byte(s)); err != nil {
t.Fatal(err)
}
}
if err := rocksdb.Flush(); err != nil {
t.Fatal(err)
}
ssts, err := rocksdb.getUserProperties()
if err != nil {
t.Fatal(err)
}
if len(ssts.Sst) != 1 {
t.Fatalf("expected 1 sstable got %d", len(ssts.Sst))
}
sst := ssts.Sst[0]
if sst.TsMin == nil || !sst.TsMin.Equal(minTimestamp) {
t.Fatalf("got min %v expected %v", sst.TsMin, minTimestamp)
}
if sst.TsMax == nil || !sst.TsMax.Equal(maxTimestamp) {
t.Fatalf("got max %v expected %v", sst.TsMax, maxTimestamp)
}
}
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
}
}
}