// compareBiogoTree walks both a biogo tree and the range tree to determine if both
// contain the same values in the same order.
func compareBiogoTree(db *client.DB, biogoTree *llrb.Tree) error {
rt := &proto.RangeTree{}
if err := db.GetProto(keys.RangeTreeRoot, rt); err != nil {
return err
}
return compareBiogoNode(db, biogoTree.Root, &rt.RootKey)
}
func countRangeReplicas(db *client.DB) (int, error) {
desc := &roachpb.RangeDescriptor{}
if err := db.GetProto(keys.RangeDescriptorKey(roachpb.KeyMin), desc); err != nil {
return 0, err
}
return len(desc.Replicas), nil
}
// getConfig retrieves the configuration for the specified key. If the
// key is empty, all configurations are returned. Otherwise, the
// leading "/" path delimiter is stripped and the configuration
// matching the remainder is retrieved. Note that this will retrieve
// the default config if "key" is equal to "/", and will list all
// configs if "key" is equal to "". The body result contains a listing
// of keys and retrieval of a config. The output format is determined
// by the request header.
func getConfig(db *client.DB, configPrefix proto.Key, config gogoproto.Message,
path string, r *http.Request) (body []byte, contentType string, err error) {
// Scan all configs if the key is empty.
if len(path) == 0 {
var rows []client.KeyValue
if rows, err = db.Scan(configPrefix, configPrefix.PrefixEnd(), maxGetResults); err != nil {
return
}
if len(rows) == maxGetResults {
log.Warningf("retrieved maximum number of results (%d); some may be missing", maxGetResults)
}
var prefixes []string
for _, row := range rows {
trimmed := bytes.TrimPrefix(row.Key, configPrefix)
prefixes = append(prefixes, url.QueryEscape(string(trimmed)))
}
// Encode the response.
body, contentType, err = util.MarshalResponse(r, prefixes, util.AllEncodings)
} else {
configkey := keys.MakeKey(configPrefix, proto.Key(path[1:]))
if err = db.GetProto(configkey, config); err != nil {
return
}
body, contentType, err = util.MarshalResponse(r, config, util.AllEncodings)
}
return
}
// compareBiogoNode compares a biogo node and a range tree node to determine if both
// contain the same values in the same order. It recursively calls itself on
// both children if they exist.
func compareBiogoNode(db *client.DB, biogoNode *llrb.Node, key *proto.Key) error {
// Retrieve the node form the range tree.
rtNode := &proto.RangeTreeNode{}
if err := db.GetProto(keys.RangeTreeNodeKey(*key), rtNode); err != nil {
return err
}
bNode := &proto.RangeTreeNode{
Key: proto.Key(biogoNode.Elem.(Key)),
ParentKey: proto.KeyMin,
Black: bool(biogoNode.Color),
}
if biogoNode.Left != nil {
leftKey := proto.Key(biogoNode.Left.Elem.(Key))
bNode.LeftKey = &leftKey
}
if biogoNode.Right != nil {
rightKey := proto.Key(biogoNode.Right.Elem.(Key))
bNode.RightKey = &rightKey
}
if err := nodesEqual(*key, *bNode, *rtNode); err != nil {
return err
}
if rtNode.LeftKey != nil {
if err := compareBiogoNode(db, biogoNode.Left, rtNode.LeftKey); err != nil {
return err
}
}
if rtNode.RightKey != nil {
if err := compareBiogoNode(db, biogoNode.Right, rtNode.RightKey); err != nil {
return err
}
}
return nil
}
// GetTableDescriptor retrieves a table descriptor directly from the KV layer.
func GetTableDescriptor(kvDB *client.DB, database string, table string) *TableDescriptor {
dbNameKey := MakeNameMetadataKey(keys.RootNamespaceID, database)
gr, err := kvDB.Get(dbNameKey)
if err != nil {
panic(err)
}
if !gr.Exists() {
panic("database missing")
}
dbDescID := ID(gr.ValueInt())
tableNameKey := MakeNameMetadataKey(dbDescID, table)
gr, err = kvDB.Get(tableNameKey)
if err != nil {
panic(err)
}
if !gr.Exists() {
panic("table missing")
}
descKey := MakeDescMetadataKey(ID(gr.ValueInt()))
desc := &Descriptor{}
if err := kvDB.GetProto(descKey, desc); err != nil {
panic("proto missing")
}
return desc.GetTable()
}
// getPermConfig fetches the permissions config for 'prefix'.
func getPermConfig(db *client.DB, prefix string) (*config.PermConfig, error) {
config := &config.PermConfig{}
if err := db.GetProto(keys.MakeKey(keys.ConfigPermissionPrefix, proto.Key(prefix)), config); err != nil {
return nil, err
}
return config, nil
}
// loadTree loads the tree root and all of its nodes. It puts all of the nodes
// into a map.
func loadTree(t *testing.T, db *client.DB) (*roachpb.RangeTree, map[string]roachpb.RangeTreeNode) {
tree := new(roachpb.RangeTree)
if err := db.GetProto(keys.RangeTreeRoot, tree); err != nil {
t.Fatal(err)
}
nodes := make(map[string]roachpb.RangeTreeNode)
if tree.RootKey != nil {
loadNodes(t, db, tree.RootKey, nodes)
}
return tree, nodes
}
// loadTree loads the tree root and all of its nodes. It puts all of the nodes
// into a map.
func loadTree(t *testing.T, db *client.DB) (storage.RangeTree, map[string]storage.RangeTreeNode) {
var tree storage.RangeTree
if err := db.GetProto(keys.RangeTreeRoot, &tree); err != nil {
t.Fatal(err)
}
nodes := make(map[string]storage.RangeTreeNode)
if tree.RootKey != nil {
loadNodes(t, db, tree.RootKey, nodes)
}
return tree, nodes
}
// treesEqual compares the expectedTree and expectedNodes to the actual range
// tree stored in the db.
func treesEqual(db *client.DB, expected testRangeTree) error {
// Compare the tree roots.
actualTree := &proto.RangeTree{}
if err := db.GetProto(keys.RangeTreeRoot, actualTree); err != nil {
return err
}
if !reflect.DeepEqual(&expected.Tree, actualTree) {
return util.Errorf("Range tree root is not as expected - expected:%+v - actual:%+v", expected.Tree, actualTree)
}
return treeNodesEqual(db, expected, expected.Tree.RootKey)
}
// setDefaultRangeMaxBytes sets the range-max-bytes value for the default zone.
func setDefaultRangeMaxBytes(t *testing.T, db *client.DB, maxBytes int64) {
zone := &proto.ZoneConfig{}
if err := db.GetProto(keys.ConfigZonePrefix, zone); err != nil {
t.Fatal(err)
}
if zone.RangeMaxBytes == maxBytes {
return
}
zone.RangeMaxBytes = maxBytes
if err := db.Put(keys.ConfigZonePrefix, zone); err != nil {
t.Fatal(err)
}
}
// loadNodes fetches a node and recursively all of its children.
func loadNodes(t *testing.T, db *client.DB, key roachpb.RKey, nodes map[string]roachpb.RangeTreeNode) {
node := new(roachpb.RangeTreeNode)
if err := db.GetProto(keys.RangeTreeNodeKey(key), node); err != nil {
t.Fatal(err)
}
nodes[node.Key.String()] = *node
if node.LeftKey != nil {
loadNodes(t, db, node.LeftKey, nodes)
}
if node.RightKey != nil {
loadNodes(t, db, node.RightKey, nodes)
}
}
// treeNodesEqual compares the expectedTree from the provided key to the actual
// nodes retrieved from the db. It recursively calls itself on both left and
// right children if they exist.
func treeNodesEqual(db *client.DB, expected testRangeTree, key proto.Key) error {
expectedNode, ok := expected.Nodes[string(key)]
if !ok {
return util.Errorf("Expected does not contain a node for %s", key)
}
actualNode := &proto.RangeTreeNode{}
if err := db.GetProto(keys.RangeTreeNodeKey(key), actualNode); err != nil {
return err
}
if err := nodesEqual(key, expectedNode, *actualNode); err != nil {
return err
}
if expectedNode.LeftKey != nil {
if err := treeNodesEqual(db, expected, *expectedNode.LeftKey); err != nil {
return err
}
}
if expectedNode.RightKey != nil {
if err := treeNodesEqual(db, expected, *expectedNode.RightKey); err != nil {
return err
}
}
return nil
}
// Run a particular schema change and run some OLTP operations in parallel, as
// soon as the schema change starts executing its backfill.
func runSchemaChangeWithOperations(
t *testing.T,
sqlDB *gosql.DB,
kvDB *client.DB,
schemaChange string,
maxValue int,
keyMultiple int,
descKey roachpb.Key,
backfillNotification chan bool,
) {
desc := &sqlbase.Descriptor{}
if err := kvDB.GetProto(descKey, desc); err != nil {
t.Fatal(err)
}
tableDesc := desc.GetTable()
// Run the schema change in a separate goroutine.
var wg sync.WaitGroup
wg.Add(1)
go func() {
start := timeutil.Now()
// Start schema change that eventually runs a backfill.
if _, err := sqlDB.Exec(schemaChange); err != nil {
t.Error(err)
}
t.Logf("schema change %s took %v", schemaChange, timeutil.Since(start))
wg.Done()
}()
// Wait until the schema change backfill starts.
<-backfillNotification
// Run a variety of operations during the backfill.
// Grabbing a schema change lease on the table will fail, disallowing
// another schema change from being simultaneously executed.
sc := csql.NewSchemaChangerForTesting(tableDesc.ID, 0, 0, *kvDB, nil)
if l, err := sc.AcquireLease(); err == nil {
t.Fatalf("schema change lease acquisition on table %d succeeded: %v", tableDesc.ID, l)
}
// Update some rows.
var updatedKeys []int
for i := 0; i < 10; i++ {
k := rand.Intn(maxValue)
v := maxValue + i + 1
if _, err := sqlDB.Exec(`UPDATE t.test SET v = $2 WHERE k = $1`, k, v); err != nil {
t.Fatal(err)
}
updatedKeys = append(updatedKeys, k)
}
// Reupdate updated values back to what they were before.
for _, k := range updatedKeys {
if _, err := sqlDB.Exec(`UPDATE t.test SET v = $2 WHERE k = $1`, k, maxValue-k); err != nil {
t.Fatal(err)
}
}
// Delete some rows.
deleteStartKey := rand.Intn(maxValue - 10)
for i := 0; i < 10; i++ {
if _, err := sqlDB.Exec(`DELETE FROM t.test WHERE k = $1`, deleteStartKey+i); err != nil {
t.Fatal(err)
}
}
// Reinsert deleted rows.
for i := 0; i < 10; i++ {
k := deleteStartKey + i
if _, err := sqlDB.Exec(`INSERT INTO t.test VALUES($1, $2)`, k, maxValue-k); err != nil {
t.Fatal(err)
}
}
// Insert some new rows.
numInserts := 10
for i := 0; i < numInserts; i++ {
if _, err := sqlDB.Exec(`INSERT INTO t.test VALUES($1, $2)`, maxValue+i+1, maxValue+i+1); err != nil {
t.Fatal(err)
}
}
wg.Wait() // for schema change to complete.
// Verify the number of keys left behind in the table to validate schema
// change operations.
tablePrefix := roachpb.Key(keys.MakeTablePrefix(uint32(tableDesc.ID)))
tableEnd := tablePrefix.PrefixEnd()
if kvs, err := kvDB.Scan(tablePrefix, tableEnd, 0); err != nil {
t.Fatal(err)
} else if e := keyMultiple * (maxValue + numInserts + 1); len(kvs) != e {
t.Fatalf("expected %d key value pairs, but got %d", e, len(kvs))
}
// Delete the rows inserted.
for i := 0; i < numInserts; i++ {
if _, err := sqlDB.Exec(`DELETE FROM t.test WHERE k = $1`, maxValue+i+1); err != nil {
t.Fatal(err)
}
}
}