// MakeZoneKey returns the key for 'id's entry in the system.zones table.
func MakeZoneKey(id ID) roachpb.Key {
	k := keys.MakeTablePrefix(uint32(ZonesTable.ID))
	k = encoding.EncodeUvarint(k, uint64(ZonesTable.PrimaryIndex.ID))
	k = encoding.EncodeUvarint(k, uint64(id))
	k = encoding.EncodeUvarint(k, uint64(ZonesTable.Columns[1].ID))
	return k
}

func TestMakeTableIndexKey(t *testing.T) {
	defer leaktest.AfterTest(t)
	key := MakeTableIndexKey(12, 345, []byte("foo"), []byte("bar"))
	expKey := MakeKey(TableDataPrefix, encoding.EncodeUvarint(nil, 12), encoding.EncodeUvarint(nil, 345), encoding.EncodeBytes(nil, []byte("foo")), encoding.EncodeBytes(nil, []byte("bar")))
	if !key.Equal(expKey) {
		t.Errorf("key %q doesn't match expected %q", key, expKey)
	}
	// Check that keys are ordered
	keys := []proto.Key{
		MakeTableIndexKey(0, 0, []byte("foo")),
		MakeTableIndexKey(0, 0, []byte("fooo")),
		MakeTableIndexKey(0, 1, []byte("bar")),
		MakeTableIndexKey(1, 0, []byte("bar")),
		MakeTableIndexKey(1, 0, []byte("bar"), []byte("foo")),
		MakeTableIndexKey(1, 1, []byte("bar"), []byte("fo")),
		MakeTableIndexKey(1, 1, []byte("bar"), []byte("foo")),
		MakeTableIndexKey(1, 2, []byte("bar")),
		MakeTableIndexKey(2, 2, []byte("ba")),
	}
	for i := 1; i < len(keys); i++ {
		if bytes.Compare(keys[i-1], keys[i]) >= 0 {
			t.Errorf("key %d >= key %d", i-1, i)
		}
	}
}

func newRangeDataIterator(d *proto.RangeDescriptor, e engine.Engine) *rangeDataIterator {
	// The first range in the keyspace starts at KeyMin, which includes the node-local
	// space. We need the original StartKey to find the range metadata, but the
	// actual data starts at LocalMax.
	dataStartKey := d.StartKey
	if d.StartKey.Equal(proto.KeyMin) {
		dataStartKey = keys.LocalMax
	}
	ri := &rangeDataIterator{
		ranges: []keyRange{
			{
				start: engine.MVCCEncodeKey(keys.MakeKey(keys.LocalRangeIDPrefix, encoding.EncodeUvarint(nil, uint64(d.RangeID)))),
				end:   engine.MVCCEncodeKey(keys.MakeKey(keys.LocalRangeIDPrefix, encoding.EncodeUvarint(nil, uint64(d.RangeID+1)))),
			},
			{
				start: engine.MVCCEncodeKey(keys.MakeKey(keys.LocalRangePrefix, encoding.EncodeBytes(nil, d.StartKey))),
				end:   engine.MVCCEncodeKey(keys.MakeKey(keys.LocalRangePrefix, encoding.EncodeBytes(nil, d.EndKey))),
			},
			{
				start: engine.MVCCEncodeKey(dataStartKey),
				end:   engine.MVCCEncodeKey(d.EndKey),
			},
		},
		iter: e.NewIterator(),
	}
	ri.iter.Seek(ri.ranges[ri.curIndex].start)
	ri.advance()
	return ri
}

func sqlKV(tableID uint32, indexID, descriptorID uint64) roachpb.KeyValue {
	k := keys.MakeTablePrefix(tableID)
	k = encoding.EncodeUvarint(k, indexID)
	k = encoding.EncodeUvarint(k, descriptorID)
	k = encoding.EncodeUvarint(k, 12345) // Column ID, but could be anything.
	return kv(k, nil)
}

// runClientScan first creates test data (and resets the benchmarking
// timer). It then performs b.N client scans in increments of numRows
// keys over all of the data, restarting at the beginning of the
// keyspace, as many times as necessary.
func runClientScan(useSSL bool, numRows, numVersions int, b *testing.B) {
	const numKeys = 100000

	s, db := setupClientBenchData(useSSL, numVersions, numKeys, b)
	defer s.Stop()

	b.SetBytes(int64(numRows * valueSize))
	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		startKeyBuf := append(make([]byte, 0, 64), []byte("key-")...)
		endKeyBuf := append(make([]byte, 0, 64), []byte("key-")...)
		for pb.Next() {
			// Choose a random key to start scan.
			keyIdx := rand.Int31n(int32(numKeys - numRows))
			startKey := roachpb.Key(encoding.EncodeUvarint(startKeyBuf, uint64(keyIdx)))
			endKey := roachpb.Key(encoding.EncodeUvarint(endKeyBuf, uint64(keyIdx)+uint64(numRows)))
			rows, err := db.Scan(startKey, endKey, int64(numRows))
			if err != nil {
				b.Fatalf("failed scan: %s", err)
			}
			if len(rows) != numRows {
				b.Fatalf("failed to scan: %d != %d", len(rows), numRows)
			}
		}
	})

	b.StopTimer()
}

// MakeDescMetadataKey returns the key for the descriptor.
func MakeDescMetadataKey(descID ID) proto.Key {
	k := MakeTablePrefix(DescriptorTable.ID)
	k = encoding.EncodeUvarint(k, uint64(DescriptorTable.PrimaryIndex.ID))
	k = encoding.EncodeUvarint(k, uint64(descID))
	k = encoding.EncodeUvarint(k, uint64(DescriptorTable.Columns[1].ID))
	return k
}

func encodeIndexKeyPrefix(tableID, indexID uint32) []byte {
	var key []byte
	key = append(key, keys.TableDataPrefix...)
	key = encoding.EncodeUvarint(key, uint64(tableID))
	key = encoding.EncodeUvarint(key, uint64(indexID))
	return key
}

func newRangeDataIterator(r *Range, e engine.Engine) *rangeDataIterator {
	r.RLock()
	startKey := r.Desc().StartKey
	if startKey.Equal(engine.KeyMin) {
		startKey = engine.KeyLocalMax
	}
	endKey := r.Desc().EndKey
	r.RUnlock()
	ri := &rangeDataIterator{
		ranges: []keyRange{
			{
				start: engine.MVCCEncodeKey(engine.MakeKey(engine.KeyLocalRangeIDPrefix, encoding.EncodeUvarint(nil, uint64(r.Desc().RaftID)))),
				end:   engine.MVCCEncodeKey(engine.MakeKey(engine.KeyLocalRangeIDPrefix, encoding.EncodeUvarint(nil, uint64(r.Desc().RaftID+1)))),
			},
			{
				start: engine.MVCCEncodeKey(engine.MakeKey(engine.KeyLocalRangeKeyPrefix, encoding.EncodeBytes(nil, startKey))),
				end:   engine.MVCCEncodeKey(engine.MakeKey(engine.KeyLocalRangeKeyPrefix, encoding.EncodeBytes(nil, endKey))),
			},
			{
				start: engine.MVCCEncodeKey(startKey),
				end:   engine.MVCCEncodeKey(endKey),
			},
		},
		iter: e.NewIterator(),
	}
	ri.iter.Seek(ri.ranges[ri.curIndex].start)
	ri.advance()
	return ri
}

// MakeDescMetadataKey returns the key for the descriptor.
func MakeDescMetadataKey(descID ID) roachpb.Key {
	k := keys.MakeTablePrefix(uint32(DescriptorTable.ID))
	k = encoding.EncodeUvarint(k, uint64(DescriptorTable.PrimaryIndex.ID))
	k = encoding.EncodeUvarint(k, uint64(descID))
	k = encoding.EncodeUvarint(k, uint64(DescriptorTable.Columns[1].ID))
	return k
}

// populateTableIndexKey populates the key passed in with the
// order encoded values forming the index key.
func populateTableIndexKey(key []byte, tableID, indexID uint32, columnValues ...[]byte) []byte {
	key = append(key, TableDataPrefix...)
	key = encoding.EncodeUvarint(key, uint64(tableID))
	key = encoding.EncodeUvarint(key, uint64(indexID))
	for _, value := range columnValues {
		key = encoding.EncodeBytes(key, value)
	}
	return key
}

// MakeColumnKey returns the key for the column in the given row.
func MakeColumnKey(rowKey []byte, colID uint32) []byte {
	key := append([]byte(nil), rowKey...)
	size := len(key)
	key = encoding.EncodeUvarint(key, uint64(colID))
	// Note that we assume that `len(key)-size` will always be encoded to a
	// single byte by EncodeUvarint. This is currently always true because the
	// varint encoding will encode 1-9 bytes.
	return encoding.EncodeUvarint(key, uint64(len(key)-size))
}

// MakeNameMetadataKey returns the key for the name. Pass name == "" in order
// to generate the prefix key to use to scan over all of the names for the
// specified parentID.
func MakeNameMetadataKey(parentID ID, name string) roachpb.Key {
	name = normalizeName(name)
	k := keys.MakeTablePrefix(uint32(NamespaceTable.ID))
	k = encoding.EncodeUvarint(k, uint64(NamespaceTable.PrimaryIndex.ID))
	k = encoding.EncodeUvarint(k, uint64(parentID))
	if name != "" {
		k = encoding.EncodeBytes(k, []byte(name))
		k = keys.MakeColumnKey(k, uint32(NamespaceTable.Columns[2].ID))
	}
	return k
}

// runMVCCScan first creates test data (and resets the benchmarking
// timer). It then performs b.N MVCCScans in increments of numRows
// keys over all of the data in the rocksdb instance, restarting at
// the beginning of the keyspace, as many times as necessary.
func runMVCCScan(numRows, numVersions int, b *testing.B) {
	// Use the same number of keys for all of the mvcc scan
	// benchmarks. Using a different number of keys per test gives
	// preferential treatment to tests with fewer keys. Note that the
	// datasets all fit in cache and the cache is pre-warmed.
	const numKeys = 100000

	rocksdb := setupMVCCScanData(numVersions, numKeys, b)
	defer rocksdb.Stop()

	prewarmCache(rocksdb)

	b.SetBytes(int64(numRows * 1024))
	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			// Choose a random key to start scan.
			keyIdx := rand.Int31n(int32(numKeys - numRows))
			startKey := proto.Key(encoding.EncodeUvarint([]byte("key-"), uint64(keyIdx)))
			walltime := int64(5 * (rand.Int31n(int32(numVersions)) + 1))
			ts := makeTS(walltime, 0)
			kvs, err := MVCCScan(rocksdb, startKey, KeyMax, int64(numRows), ts, nil)
			if err != nil {
				b.Fatalf("failed scan: %s", err)
			}
			if len(kvs) != numRows {
				b.Fatalf("failed to scan: %d != %d", len(kvs), numRows)
			}
		}
	})

	b.StopTimer()
}

// MakeNameMetadataKey returns the key for the namespace.
func MakeNameMetadataKey(parentID uint32, name string) proto.Key {
	k := make([]byte, 0, len(NameMetadataPrefix)+encoding.MaxUvarintSize+len(name))
	k = append(k, NameMetadataPrefix...)
	k = encoding.EncodeUvarint(k, uint64(parentID))
	k = append(k, name...)
	return k
}

// encodeTableKey encodes a single element of a table key, appending the
// encoded value to b.
func encodeTableKey(b []byte, v reflect.Value) ([]byte, error) {
	switch t := v.Interface().(type) {
	case []byte:
		return roachencoding.EncodeBytes(b, t), nil
	case string:
		return roachencoding.EncodeBytes(b, []byte(t)), nil
	}

	switch v.Kind() {
	case reflect.Bool:
		if v.Bool() {
			return roachencoding.EncodeVarint(b, 1), nil
		}
		return roachencoding.EncodeVarint(b, 0), nil

	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return roachencoding.EncodeVarint(b, v.Int()), nil

	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
		return roachencoding.EncodeUvarint(b, v.Uint()), nil

	case reflect.Float32, reflect.Float64:
		return roachencoding.EncodeNumericFloat(b, v.Float()), nil

	case reflect.String:
		return roachencoding.EncodeBytes(b, []byte(v.String())), nil
	}

	return nil, fmt.Errorf("unable to encode key: %s", v)
}

func TestMakeTableDataKey(t *testing.T) {
	defer leaktest.AfterTest(t)
	key := MakeTableDataKey(12, 345, 6, []byte("foo"), []byte("bar"))
	// Expected key is the TableIndexKey + ColumnID
	expKey := MakeKey(MakeTableIndexKey(12, 345, []byte("foo"), []byte("bar")), encoding.EncodeUvarint(nil, 6))
	if !key.Equal(expKey) {
		t.Errorf("key %q doesn't match expected %q", key, expKey)
	}
	// Check that keys are ordered
	keys := []proto.Key{
		// Data-key follows Index key order
		MakeTableDataKey(0, 0, 0, []byte("foo")),
		MakeTableDataKey(0, 0, 8, []byte("fooo")),
		MakeTableDataKey(0, 1, 10, []byte("bar")),
		MakeTableDataKey(1, 0, 3, []byte("bar")),
		MakeTableDataKey(1, 0, 5, []byte("bar"), []byte("foo")),
		MakeTableDataKey(1, 1, 7, []byte("bar"), []byte("fo")),
		MakeTableDataKey(1, 1, 4, []byte("bar"), []byte("foo")),
		MakeTableDataKey(1, 2, 89, []byte("bar")),
		MakeTableDataKey(2, 2, 23, []byte("ba")),
		// For the same Index key, Data-key follows column ID order.
		MakeTableDataKey(2, 2, 0, []byte("bar"), []byte("foo")),
		MakeTableDataKey(2, 2, 7, []byte("bar"), []byte("foo")),
		MakeTableDataKey(2, 2, 23, []byte("bar"), []byte("foo")),
		MakeTableDataKey(2, 2, 45, []byte("bar"), []byte("foo")),
	}

	for i := 1; i < len(keys); i++ {
		if bytes.Compare(keys[i-1], keys[i]) >= 0 {
			t.Errorf("key %d >= key %d", i-1, i)
		}
	}

}

func runMVCCBatchPut(valueSize, batchSize int, b *testing.B) {
	rng := util.NewPseudoRand()
	value := proto.Value{Bytes: util.RandBytes(rng, valueSize)}
	keyBuf := append(make([]byte, 0, 64), []byte("key-")...)

	rocksdb := NewInMem(proto.Attributes{Attrs: []string{"ssd"}}, testCacheSize)
	defer rocksdb.Stop()

	b.SetBytes(int64(valueSize))
	b.ResetTimer()

	for i := 0; i < b.N; i += batchSize {
		end := i + batchSize
		if end > b.N {
			end = b.N
		}

		batch := rocksdb.NewBatch()

		for j := i; j < end; j++ {
			key := proto.Key(encoding.EncodeUvarint(keyBuf[0:4], uint64(j)))
			ts := makeTS(time.Now().UnixNano(), 0)
			if err := MVCCPut(batch, nil, key, ts, value, nil); err != nil {
				b.Fatalf("failed put: %s", err)
			}
		}

		if err := batch.Commit(); err != nil {
			b.Fatal(err)
		}
	}

	b.StopTimer()
}

// runMVCCGet first creates test data (and resets the benchmarking
// timer). It then performs b.N MVCCGets.
func runMVCCGet(numVersions, valueSize int, b *testing.B) {
	const overhead = 48          // Per key/value overhead (empirically determined)
	const targetSize = 512 << 20 // 512 MB
	// Adjust the number of keys so that each test has approximately the same
	// amount of data.
	numKeys := targetSize / ((overhead + valueSize) * (1 + (numVersions-1)/2))

	rocksdb, stopper := setupMVCCData(numVersions, numKeys, valueSize, b)
	defer stopper.Stop()

	b.SetBytes(int64(valueSize))
	b.ResetTimer()

	keyBuf := append(make([]byte, 0, 64), []byte("key-")...)
	for i := 0; i < b.N; i++ {
		// Choose a random key to retrieve.
		keyIdx := rand.Int31n(int32(numKeys))
		key := roachpb.Key(encoding.EncodeUvarint(keyBuf[:4], uint64(keyIdx)))
		walltime := int64(5 * (rand.Int31n(int32(numVersions)) + 1))
		ts := makeTS(walltime, 0)
		if v, _, err := MVCCGet(rocksdb, key, ts, true, nil); err != nil {
			b.Fatalf("failed get: %s", err)
		} else if v == nil {
			b.Fatalf("failed get (key not found): %[email protected]%d", keyIdx, walltime)
		} else if valueBytes, err := v.GetBytes(); err != nil {
			b.Fatal(err)
		} else if len(valueBytes) != valueSize {
			b.Fatalf("unexpected value size: %d", len(valueBytes))
		}
	}

	b.StopTimer()
}

// runMVCCGet first creates test data (and resets the benchmarking
// timer). It then performs b.N MVCCGets.
func runMVCCGet(numVersions int, b *testing.B) {
	// Use the same number of keys for all of the mvcc get
	// benchmarks. Using a different number of keys per test gives
	// preferential treatment to tests with fewer keys. Note that the
	// datasets all fit in cache and the cache is pre-warmed.
	const numKeys = 100000

	rocksdb := setupMVCCScanData(numVersions, numKeys, b)
	defer rocksdb.Close()

	prewarmCache(rocksdb)

	b.SetBytes(1024)
	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		keyBuf := append(make([]byte, 0, 64), []byte("key-")...)
		for pb.Next() {
			// Choose a random key to retrieve.
			keyIdx := rand.Int31n(int32(numKeys))
			key := proto.Key(encoding.EncodeUvarint(keyBuf[0:4], uint64(keyIdx)))
			walltime := int64(5 * (rand.Int31n(int32(numVersions)) + 1))
			ts := makeTS(walltime, 0)
			if v, _, err := MVCCGet(rocksdb, key, ts, true, nil); err != nil {
				b.Fatalf("failed get: %s", err)
			} else if len(v.Bytes) != 1024 {
				b.Fatalf("unexpected value size: %d", len(v.Bytes))
			}
		}
	})

	b.StopTimer()
}

// runMVCCScan first creates test data (and resets the benchmarking
// timer). It then performs b.N MVCCScans in increments of numRows
// keys over all of the data in the rocksdb instance, restarting at
// the beginning of the keyspace, as many times as necessary.
func runMVCCScan(numRows, numVersions, valueSize int, b *testing.B) {
	// Use the same number of keys for all of the mvcc scan
	// benchmarks. Using a different number of keys per test gives
	// preferential treatment to tests with fewer keys. Note that the
	// datasets all fit in cache and the cache is pre-warmed.
	const numKeys = 100000

	rocksdb, stopper := setupMVCCData(numVersions, numKeys, valueSize, b)
	defer stopper.Stop()

	b.SetBytes(int64(numRows * valueSize))
	b.ResetTimer()

	keyBuf := append(make([]byte, 0, 64), []byte("key-")...)
	for i := 0; i < b.N; i++ {
		// Choose a random key to start scan.
		keyIdx := rand.Int31n(int32(numKeys - numRows))
		startKey := roachpb.Key(encoding.EncodeUvarint(keyBuf[:4], uint64(keyIdx)))
		walltime := int64(5 * (rand.Int31n(int32(numVersions)) + 1))
		ts := makeTS(walltime, 0)
		kvs, _, err := MVCCScan(rocksdb, startKey, keyMax, int64(numRows), ts, true, nil)
		if err != nil {
			b.Fatalf("failed scan: %s", err)
		}
		if len(kvs) != numRows {
			b.Fatalf("failed to scan: %d != %d", len(kvs), numRows)
		}
	}

	b.StopTimer()
}