Ejemplo n.º 1
0
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
		}
	}
}
Ejemplo n.º 2
0
// 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)
}