micro/registry/mdns_registry.go
2019-02-01 13:41:11 +00:00

345 lines
6.0 KiB
Go

// Package mdns is a multicast dns registry
package registry
import (
"context"
"net"
"strings"
"sync"
"time"
"github.com/micro/mdns"
hash "github.com/mitchellh/hashstructure"
)
type mdnsTxt struct {
Service string
Version string
Endpoints []*Endpoint
Metadata map[string]string
}
type mdnsEntry struct {
hash uint64
id string
node *mdns.Server
}
type mdnsRegistry struct {
opts Options
sync.Mutex
services map[string][]*mdnsEntry
}
func newRegistry(opts ...Option) Registry {
options := Options{
Timeout: time.Millisecond * 100,
}
return &mdnsRegistry{
opts: options,
services: make(map[string][]*mdnsEntry),
}
}
func (m *mdnsRegistry) Init(opts ...Option) error {
for _, o := range opts {
o(&m.opts)
}
return nil
}
func (m *mdnsRegistry) Options() Options {
return m.opts
}
func (m *mdnsRegistry) Register(service *Service, opts ...RegisterOption) error {
m.Lock()
defer m.Unlock()
entries, ok := m.services[service.Name]
// first entry, create wildcard used for list queries
if !ok {
s, err := mdns.NewMDNSService(
service.Name,
"_services",
"",
"",
9999,
[]net.IP{net.ParseIP("0.0.0.0")},
nil,
)
if err != nil {
return err
}
srv, err := mdns.NewServer(&mdns.Config{Zone: &mdns.DNSSDService{s}})
if err != nil {
return err
}
// append the wildcard entry
entries = append(entries, &mdnsEntry{id: "*", node: srv})
}
var gerr error
for _, node := range service.Nodes {
// create hash of service; uint64
h, err := hash.Hash(node, nil)
if err != nil {
gerr = err
continue
}
var seen bool
var e *mdnsEntry
for _, entry := range entries {
if node.Id == entry.id {
seen = true
e = entry
break
}
}
// already registered, continue
if seen && e.hash == h {
continue
// hash doesn't match, shutdown
} else if seen {
e.node.Shutdown()
// doesn't exist
} else {
e = &mdnsEntry{hash: h}
}
txt, err := encode(&mdnsTxt{
Service: service.Name,
Version: service.Version,
Endpoints: service.Endpoints,
Metadata: node.Metadata,
})
if err != nil {
gerr = err
continue
}
// we got here, new node
s, err := mdns.NewMDNSService(
node.Id,
service.Name,
"",
"",
node.Port,
[]net.IP{net.ParseIP(node.Address)},
txt,
)
if err != nil {
gerr = err
continue
}
srv, err := mdns.NewServer(&mdns.Config{Zone: s})
if err != nil {
gerr = err
continue
}
e.id = node.Id
e.node = srv
entries = append(entries, e)
}
// save
m.services[service.Name] = entries
return gerr
}
func (m *mdnsRegistry) Deregister(service *Service) error {
m.Lock()
defer m.Unlock()
var newEntries []*mdnsEntry
// loop existing entries, check if any match, shutdown those that do
for _, entry := range m.services[service.Name] {
var remove bool
for _, node := range service.Nodes {
if node.Id == entry.id {
entry.node.Shutdown()
remove = true
break
}
}
// keep it?
if !remove {
newEntries = append(newEntries, entry)
}
}
// last entry is the wildcard for list queries. Remove it.
if len(newEntries) == 1 && newEntries[0].id == "*" {
newEntries[0].node.Shutdown()
delete(m.services, service.Name)
} else {
m.services[service.Name] = newEntries
}
return nil
}
func (m *mdnsRegistry) GetService(service string) ([]*Service, error) {
serviceMap := make(map[string]*Service)
entries := make(chan *mdns.ServiceEntry, 10)
done := make(chan bool)
p := mdns.DefaultParams(service)
// set context with timeout
p.Context, _ = context.WithTimeout(context.Background(), m.opts.Timeout)
// set entries channel
p.Entries = entries
go func() {
for {
select {
case e := <-entries:
// list record so skip
if p.Service == "_services" {
continue
}
if e.TTL == 0 {
continue
}
txt, err := decode(e.InfoFields)
if err != nil {
continue
}
if txt.Service != service {
continue
}
s, ok := serviceMap[txt.Version]
if !ok {
s = &Service{
Name: txt.Service,
Version: txt.Version,
Endpoints: txt.Endpoints,
}
}
s.Nodes = append(s.Nodes, &Node{
Id: strings.TrimSuffix(e.Name, "."+p.Service+"."+p.Domain+"."),
Address: e.AddrV4.String(),
Port: e.Port,
Metadata: txt.Metadata,
})
serviceMap[txt.Version] = s
case <-p.Context.Done():
close(done)
return
}
}
}()
// execute the query
if err := mdns.Query(p); err != nil {
return nil, err
}
// wait for completion
<-done
// create list and return
var services []*Service
for _, service := range serviceMap {
services = append(services, service)
}
return services, nil
}
func (m *mdnsRegistry) ListServices() ([]*Service, error) {
serviceMap := make(map[string]bool)
entries := make(chan *mdns.ServiceEntry, 10)
done := make(chan bool)
p := mdns.DefaultParams("_services")
// set context with timeout
p.Context, _ = context.WithTimeout(context.Background(), m.opts.Timeout)
// set entries channel
p.Entries = entries
var services []*Service
go func() {
for {
select {
case e := <-entries:
if e.TTL == 0 {
continue
}
name := strings.TrimSuffix(e.Name, "."+p.Service+"."+p.Domain+".")
if !serviceMap[name] {
serviceMap[name] = true
services = append(services, &Service{Name: name})
}
case <-p.Context.Done():
close(done)
return
}
}
}()
// execute query
if err := mdns.Query(p); err != nil {
return nil, err
}
// wait till done
<-done
return services, nil
}
func (m *mdnsRegistry) Watch(opts ...WatchOption) (Watcher, error) {
var wo WatchOptions
for _, o := range opts {
o(&wo)
}
md := &mdnsWatcher{
wo: wo,
ch: make(chan *mdns.ServiceEntry, 32),
exit: make(chan struct{}),
}
go func() {
if err := mdns.Listen(md.ch, md.exit); err != nil {
md.Stop()
}
}()
return md, nil
}
func (m *mdnsRegistry) String() string {
return "mdns"
}
// NewRegistry returns a new default registry which is mdns
func NewRegistry(opts ...Option) Registry {
return newRegistry(opts...)
}