Golang client.DB类代码示例

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

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

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

// allocateNodeID increments the node id generator key to allocate
// a new, unique node id.
func allocateNodeID(db *client.DB) (roachpb.NodeID, error) {
	r, err := db.Inc(keys.NodeIDGenerator, 1)
	if err != nil {
		return 0, util.Errorf("unable to allocate node ID: %s", err)
	}
	return roachpb.NodeID(r.ValueInt()), nil
}

func (p *planner) releaseLeases(db client.DB) {
	if p.leases != nil {
		for _, lease := range p.leases {
			if err := p.leaseMgr.Release(lease); err != nil {
				log.Warning(err)
			}
		}
		p.leases = nil
	}

	// TODO(pmattis): This is a hack. Remove when schema change operations work
	// properly.
	if p.modifiedSchemas != nil {
		for _, d := range p.modifiedSchemas {
			var lease *LeaseState
			err := db.Txn(func(txn *client.Txn) error {
				var err error
				lease, err = p.leaseMgr.Acquire(txn, d.id, d.version)
				return err
			})
			if err != nil {
				log.Warning(err)
				continue
			}
			if err := p.leaseMgr.Release(lease); err != nil {
				log.Warning(err)
			}
		}
		p.modifiedSchemas = nil
	}
}

// allocateStoreIDs increments the store id generator key for the
// specified node to allocate "inc" new, unique store ids. The
// first ID in a contiguous range is returned on success.
func allocateStoreIDs(nodeID roachpb.NodeID, inc int64, db *client.DB) (roachpb.StoreID, error) {
	r, err := db.Inc(keys.StoreIDGenerator, inc)
	if err != nil {
		return 0, util.Errorf("unable to allocate %d store IDs for node %d: %s", inc, nodeID, err)
	}
	return roachpb.StoreID(r.ValueInt() - inc + 1), nil
}

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

func split(db *client.DB) {
	for i := 1; i < 10; i++ {
		fmt.Printf("split key: %v\n", i)
		err := db.AdminSplit([]byte(fmt.Sprintf("%d", i)))
		if err != nil {
			panic(fmt.Sprintf("split fail. key: %v, err: %v\n", i, err))
		}
	}
}

// deleteConfig removes the config specified by key.
func deleteConfig(db *client.DB, configPrefix proto.Key, path string, r *http.Request) error {
	if len(path) == 0 {
		return util.Errorf("no path specified for config Delete")
	}
	if path == "/" {
		return util.Errorf("the default configuration cannot be deleted")
	}
	configKey := keys.MakeKey(configPrefix, proto.Key(path[1:]))
	return db.Del(configKey)
}

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

// getRangeKeys returns the end keys of all ranges.
func getRangeKeys(db *client.DB) ([]roachpb.Key, error) {
	rows, err := db.Scan(keys.Meta2Prefix, keys.MetaMax, 0)
	if err != nil {
		return nil, err
	}
	ret := make([]roachpb.Key, len(rows), len(rows))
	for i := 0; i < len(rows); i++ {
		ret[i] = bytes.TrimPrefix(rows[i].Key, keys.Meta2Prefix)
	}
	return ret, nil
}

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

func insert(r *rand.Rand, db *client.DB, wg *sync.WaitGroup) {
	wg.Add(1)
	for i := 0; i < 1000; i++ {
		key := getKey(r)
		value := getValue(1024 * 10)
		err := db.Put(key, value)
		if err != nil {
			fmt.Printf("put fail. err: %v, key: %v\n", err, key)
		}
	}
	wg.Done()
}