322 lines
7.5 KiB
Go
322 lines
7.5 KiB
Go
package server
|
|
|
|
import (
|
|
"fmt"
|
|
"net/url"
|
|
"path"
|
|
"path/filepath"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/coreos/etcd/log"
|
|
"github.com/coreos/etcd/store"
|
|
)
|
|
|
|
// The location of the peer URL data.
|
|
const RegistryPeerKey = "/_etcd/machines"
|
|
|
|
// The location of the proxy URL data.
|
|
const RegistryProxyKey = "/_etcd/proxies"
|
|
|
|
// The Registry stores URL information for nodes.
|
|
type Registry struct {
|
|
sync.Mutex
|
|
store store.Store
|
|
peers map[string]*node
|
|
proxies map[string]*node
|
|
}
|
|
|
|
// The internal storage format of the registry.
|
|
type node struct {
|
|
peerVersion string
|
|
peerURL string
|
|
url string
|
|
}
|
|
|
|
// Creates a new Registry.
|
|
func NewRegistry(s store.Store) *Registry {
|
|
return &Registry{
|
|
store: s,
|
|
peers: make(map[string]*node),
|
|
proxies: make(map[string]*node),
|
|
}
|
|
}
|
|
|
|
// Peers returns a list of cached peer names.
|
|
func (r *Registry) Peers() []string {
|
|
r.Lock()
|
|
defer r.Unlock()
|
|
|
|
names := make([]string, 0, len(r.peers))
|
|
for name := range r.peers {
|
|
names = append(names, name)
|
|
}
|
|
sort.Sort(sort.StringSlice(names))
|
|
return names
|
|
}
|
|
|
|
// Proxies returns a list of cached proxy names.
|
|
func (r *Registry) Proxies() []string {
|
|
r.Lock()
|
|
defer r.Unlock()
|
|
|
|
names := make([]string, 0, len(r.proxies))
|
|
for name := range r.proxies {
|
|
names = append(names, name)
|
|
}
|
|
sort.Sort(sort.StringSlice(names))
|
|
return names
|
|
}
|
|
|
|
// RegisterPeer adds a peer to the registry.
|
|
func (r *Registry) RegisterPeer(name string, peerURL string, machURL string) error {
|
|
if err := r.register(RegistryPeerKey, name, peerURL, machURL); err != nil {
|
|
return err
|
|
}
|
|
|
|
r.Lock()
|
|
defer r.Unlock()
|
|
r.peers[name] = r.load(RegistryPeerKey, name)
|
|
return nil
|
|
}
|
|
|
|
// RegisterProxy adds a proxy to the registry.
|
|
func (r *Registry) RegisterProxy(name string, peerURL string, machURL string) error {
|
|
if err := r.register(RegistryProxyKey, name, peerURL, machURL); err != nil {
|
|
return err
|
|
}
|
|
|
|
r.Lock()
|
|
defer r.Unlock()
|
|
r.proxies[name] = r.load(RegistryProxyKey, name)
|
|
return nil
|
|
}
|
|
|
|
func (r *Registry) register(key, name string, peerURL string, machURL string) error {
|
|
// Write data to store.
|
|
v := url.Values{}
|
|
v.Set("raft", peerURL)
|
|
v.Set("etcd", machURL)
|
|
_, err := r.store.Create(path.Join(key, name), false, v.Encode(), false, store.Permanent)
|
|
log.Debugf("Register: %s", name)
|
|
return err
|
|
}
|
|
|
|
// UnregisterPeer removes a peer from the registry.
|
|
func (r *Registry) UnregisterPeer(name string) error {
|
|
return r.unregister(RegistryPeerKey, name)
|
|
}
|
|
|
|
// UnregisterProxy removes a proxy from the registry.
|
|
func (r *Registry) UnregisterProxy(name string) error {
|
|
return r.unregister(RegistryProxyKey, name)
|
|
}
|
|
|
|
func (r *Registry) unregister(key, name string) error {
|
|
// Remove the key from the store.
|
|
_, err := r.store.Delete(path.Join(key, name), false, false)
|
|
log.Debugf("Unregister: %s", name)
|
|
return err
|
|
}
|
|
|
|
// PeerCount returns the number of peers in the cluster.
|
|
func (r *Registry) PeerCount() int {
|
|
return r.count(RegistryPeerKey)
|
|
}
|
|
|
|
// ProxyCount returns the number of proxies in the cluster.
|
|
func (r *Registry) ProxyCount() int {
|
|
return r.count(RegistryProxyKey)
|
|
}
|
|
|
|
// Returns the number of nodes in the cluster.
|
|
func (r *Registry) count(key string) int {
|
|
e, err := r.store.Get(key, false, false)
|
|
if err != nil {
|
|
return 0
|
|
}
|
|
return len(e.Node.Nodes)
|
|
}
|
|
|
|
// PeerExists checks if a peer with the given name exists.
|
|
func (r *Registry) PeerExists(name string) bool {
|
|
return r.exists(RegistryPeerKey, name)
|
|
}
|
|
|
|
// ProxyExists checks if a proxy with the given name exists.
|
|
func (r *Registry) ProxyExists(name string) bool {
|
|
return r.exists(RegistryProxyKey, name)
|
|
}
|
|
|
|
func (r *Registry) exists(key, name string) bool {
|
|
e, err := r.store.Get(path.Join(key, name), false, false)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
return (e.Node != nil)
|
|
}
|
|
|
|
// Retrieves the client URL for a given node by name.
|
|
func (r *Registry) ClientURL(name string) (string, bool) {
|
|
r.Lock()
|
|
defer r.Unlock()
|
|
return r.clientURL(RegistryPeerKey, name)
|
|
}
|
|
|
|
func (r *Registry) clientURL(key, name string) (string, bool) {
|
|
if r.peers[name] == nil {
|
|
if node := r.load(key, name); node != nil {
|
|
r.peers[name] = node
|
|
}
|
|
}
|
|
|
|
if node := r.peers[name]; node != nil {
|
|
return node.url, true
|
|
}
|
|
|
|
return "", false
|
|
}
|
|
|
|
// TODO(yichengq): have all of the code use a full URL with scheme
|
|
// and remove this method
|
|
// PeerHost retrieves the host part of peer URL for a given node by name.
|
|
func (r *Registry) PeerHost(name string) (string, bool) {
|
|
rawurl, ok := r.PeerURL(name)
|
|
if ok {
|
|
u, _ := url.Parse(rawurl)
|
|
return u.Host, ok
|
|
}
|
|
return rawurl, ok
|
|
}
|
|
|
|
// Retrieves the peer URL for a given node by name.
|
|
func (r *Registry) PeerURL(name string) (string, bool) {
|
|
r.Lock()
|
|
defer r.Unlock()
|
|
return r.peerURL(RegistryPeerKey, name)
|
|
}
|
|
|
|
func (r *Registry) peerURL(key, name string) (string, bool) {
|
|
if r.peers[name] == nil {
|
|
if node := r.load(key, name); node != nil {
|
|
r.peers[name] = node
|
|
}
|
|
}
|
|
|
|
if node := r.peers[name]; node != nil {
|
|
return node.peerURL, true
|
|
}
|
|
|
|
return "", false
|
|
}
|
|
|
|
// Retrieves the client URL for a given proxy by name.
|
|
func (r *Registry) ProxyClientURL(name string) (string, bool) {
|
|
r.Lock()
|
|
defer r.Unlock()
|
|
return r.proxyClientURL(RegistryProxyKey, name)
|
|
}
|
|
|
|
func (r *Registry) proxyClientURL(key, name string) (string, bool) {
|
|
if r.proxies[name] == nil {
|
|
if node := r.load(key, name); node != nil {
|
|
r.proxies[name] = node
|
|
}
|
|
}
|
|
if node := r.proxies[name]; node != nil {
|
|
return node.url, true
|
|
}
|
|
return "", false
|
|
}
|
|
|
|
// Retrieves the peer URL for a given proxy by name.
|
|
func (r *Registry) ProxyPeerURL(name string) (string, bool) {
|
|
r.Lock()
|
|
defer r.Unlock()
|
|
return r.proxyPeerURL(RegistryProxyKey, name)
|
|
}
|
|
|
|
func (r *Registry) proxyPeerURL(key, name string) (string, bool) {
|
|
if r.proxies[name] == nil {
|
|
if node := r.load(key, name); node != nil {
|
|
r.proxies[name] = node
|
|
}
|
|
}
|
|
if node := r.proxies[name]; node != nil {
|
|
return node.peerURL, true
|
|
}
|
|
return "", false
|
|
}
|
|
|
|
// Retrieves the Client URLs for all nodes.
|
|
func (r *Registry) ClientURLs(leaderName, selfName string) []string {
|
|
return r.urls(RegistryPeerKey, leaderName, selfName, r.clientURL)
|
|
}
|
|
|
|
// Retrieves the Peer URLs for all nodes.
|
|
func (r *Registry) PeerURLs(leaderName, selfName string) []string {
|
|
return r.urls(RegistryPeerKey, leaderName, selfName, r.peerURL)
|
|
}
|
|
|
|
// Retrieves the URLs for all nodes using url function.
|
|
func (r *Registry) urls(key, leaderName, selfName string, url func(key, name string) (string, bool)) []string {
|
|
r.Lock()
|
|
defer r.Unlock()
|
|
|
|
// Build list including the leader and self.
|
|
urls := make([]string, 0)
|
|
if url, _ := url(key, leaderName); len(url) > 0 {
|
|
urls = append(urls, url)
|
|
}
|
|
|
|
// Retrieve a list of all nodes.
|
|
if e, _ := r.store.Get(key, false, false); e != nil {
|
|
// Lookup the URL for each one.
|
|
for _, pair := range e.Node.Nodes {
|
|
_, name := filepath.Split(pair.Key)
|
|
if url, _ := url(key, name); len(url) > 0 && name != leaderName {
|
|
urls = append(urls, url)
|
|
}
|
|
}
|
|
}
|
|
|
|
log.Infof("URLs: %s: %s / %s (%s)", key, leaderName, selfName, strings.Join(urls, ","))
|
|
|
|
return urls
|
|
}
|
|
|
|
// Removes a node from the cache.
|
|
func (r *Registry) Invalidate(name string) {
|
|
r.Lock()
|
|
defer r.Unlock()
|
|
|
|
delete(r.peers, name)
|
|
delete(r.proxies, name)
|
|
}
|
|
|
|
// Loads the given node by name from the store into the cache.
|
|
func (r *Registry) load(key, name string) *node {
|
|
if name == "" {
|
|
return nil
|
|
}
|
|
|
|
// Retrieve from store.
|
|
e, err := r.store.Get(path.Join(key, name), false, false)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
|
|
// Parse as a query string.
|
|
m, err := url.ParseQuery(*e.Node.Value)
|
|
if err != nil {
|
|
panic(fmt.Sprintf("Failed to parse peers entry: %s", name))
|
|
}
|
|
|
|
// Create node.
|
|
return &node{
|
|
url: m["etcd"][0],
|
|
peerURL: m["raft"][0],
|
|
}
|
|
}
|