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) }