package registry import ( "encoding/json" "errors" "sync" consul "github.com/hashicorp/consul/api" ) type consulRegistry struct { Address string Client *consul.Client mtx sync.RWMutex services map[string]*Service } func encodeEndpoints(en []*Endpoint) []string { var tags []string for _, e := range en { if b, err := json.Marshal(e); err == nil { tags = append(tags, "e="+string(b)) } } return tags } func decodeEndpoints(tags []string) []*Endpoint { var en []*Endpoint for _, tag := range tags { if len(tag) == 0 || tag[0] != 'e' { continue } var e *Endpoint if err := json.Unmarshal([]byte(tag[2:]), &e); err == nil { en = append(en, e) } } return en } func encodeMetadata(md map[string]string) []string { var tags []string for k, v := range md { if b, err := json.Marshal(map[string]string{ k: v, }); err == nil { tags = append(tags, "t="+string(b)) } } return tags } func decodeMetadata(tags []string) map[string]string { md := make(map[string]string) for _, tag := range tags { if len(tag) == 0 || tag[0] != 't' { continue } var kv map[string]string if err := json.Unmarshal([]byte(tag[2:]), &kv); err == nil { for k, v := range kv { md[k] = v } } } return md } func newConsulRegistry(addrs []string, opts ...Option) Registry { config := consul.DefaultConfig() if len(addrs) > 0 { config.Address = addrs[0] } client, _ := consul.NewClient(config) cr := &consulRegistry{ Address: config.Address, Client: client, services: make(map[string]*Service), } return cr } func (c *consulRegistry) Deregister(s *Service) error { 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, }, nil) return err } func (c *consulRegistry) Register(s *Service) error { if len(s.Nodes) == 0 { return errors.New("Require at least one node") } node := s.Nodes[0] tags := encodeMetadata(node.Metadata) tags = append(tags, encodeEndpoints(s.Endpoints)...) _, 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, Tags: tags, }, }, nil) return err } func (c *consulRegistry) GetService(name string) (*Service, error) { c.mtx.RLock() service, ok := c.services[name] c.mtx.RUnlock() if ok { return service, nil } rsp, _, err := c.Client.Catalog().Service(name, "", nil) if err != nil { return nil, err } cs := &Service{} for _, s := range rsp { if s.ServiceName != name { continue } cs.Endpoints = decodeEndpoints(s.ServiceTags) cs.Name = s.ServiceName cs.Nodes = append(cs.Nodes, &Node{ Id: s.ServiceID, Address: s.Address, Port: s.ServicePort, Metadata: decodeMetadata(s.ServiceTags), }) } 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, &Service{Name: service}) } return services, nil } func (c *consulRegistry) Watch() (Watcher, error) { return newConsulWatcher(c) }