micro/registry/consul_registry.go

157 lines
2.8 KiB
Go
Raw Normal View History

2015-01-14 02:31:27 +03:00
package registry
import (
"errors"
2015-02-15 02:00:47 +03:00
"sync"
2015-01-14 02:31:27 +03:00
consul "github.com/hashicorp/consul/api"
2015-01-14 02:31:27 +03:00
)
type consulRegistry struct {
2015-02-15 02:00:47 +03:00
Address string
Client *consul.Client
2015-01-14 02:31:27 +03:00
2015-02-15 02:00:47 +03:00
mtx sync.RWMutex
services map[string]Service
}
2015-01-14 02:31:27 +03:00
func newConsulRegistry(addrs []string, opts ...Option) Registry {
config := consul.DefaultConfig()
client, _ := consul.NewClient(config)
if len(addrs) > 0 {
config.Address = addrs[0]
}
cr := &consulRegistry{
Address: config.Address,
Client: client,
services: make(map[string]Service),
}
cr.Watch()
return cr
}
func (c *consulRegistry) Deregister(s Service) error {
2015-01-14 02:31:27 +03:00
if len(s.Nodes()) == 0 {
return errors.New("Require at least one node")
}
node := s.Nodes()[0]
_, err := c.Client.Catalog().Deregister(&consul.CatalogDeregistration{
Node: node.Id(),
Address: node.Address(),
ServiceID: node.Id(),
}, nil)
return err
}
func (c *consulRegistry) Register(s Service) error {
2015-01-14 02:31:27 +03:00
if len(s.Nodes()) == 0 {
return errors.New("Require at least one node")
}
node := s.Nodes()[0]
_, err := c.Client.Catalog().Register(&consul.CatalogRegistration{
Node: node.Id(),
Address: node.Address(),
Service: &consul.AgentService{
ID: node.Id(),
Service: s.Name(),
Port: node.Port(),
},
}, nil)
return err
}
func (c *consulRegistry) GetService(name string) (Service, error) {
2015-02-15 02:00:47 +03:00
c.mtx.RLock()
service, ok := c.services[name]
c.mtx.RUnlock()
if ok {
return service, nil
}
2015-01-14 02:31:27 +03:00
rsp, _, err := c.Client.Catalog().Service(name, "", nil)
if err != nil {
return nil, err
}
cs := &consulService{}
2015-01-14 02:31:27 +03:00
for _, s := range rsp {
if s.ServiceName != name {
continue
}
cs.ServiceName = s.ServiceName
cs.ServiceNodes = append(cs.ServiceNodes, &consulNode{
2015-01-14 02:31:27 +03:00
Node: s.Node,
NodeId: s.ServiceID,
NodeAddress: s.Address,
NodePort: s.ServicePort,
})
}
return cs, nil
}
func (c *consulRegistry) ListServices() ([]Service, error) {
c.mtx.RLock()
serviceMap := c.services
c.mtx.RUnlock()
var services []Service
if len(serviceMap) > 0 {
for _, service := range services {
services = append(services, service)
}
return services, nil
}
rsp, _, err := c.Client.Catalog().Services(&consul.QueryOptions{})
if err != nil {
return nil, err
}
for service, _ := range rsp {
services = append(services, &consulService{ServiceName: service})
}
return services, nil
}
func (c *consulRegistry) NewService(name string, nodes ...Node) Service {
var snodes []*consulNode
2015-01-14 02:31:27 +03:00
for _, node := range nodes {
if n, ok := node.(*consulNode); ok {
2015-01-14 02:31:27 +03:00
snodes = append(snodes, n)
}
}
return &consulService{
2015-01-14 02:31:27 +03:00
ServiceName: name,
ServiceNodes: snodes,
}
}
func (c *consulRegistry) NewNode(id, address string, port int) Node {
return &consulNode{
2015-01-14 02:31:27 +03:00
Node: id,
NodeId: id,
NodeAddress: address,
NodePort: port,
}
}
func (c *consulRegistry) Watch() {
newConsulWatcher(c)
2015-01-14 02:31:27 +03:00
}